iroh_netbench/
initiator.rs

1//! Initiator-side benchmark lifecycle.
2
3use std::{
4    collections::{HashMap, HashSet, VecDeque},
5    sync::{
6        Arc, Mutex,
7        atomic::{AtomicU64, Ordering},
8    },
9    time::Duration,
10};
11
12use tokio::{
13    sync::mpsc,
14    task::{AbortHandle, JoinHandle, JoinSet},
15    time::{Instant, MissedTickBehavior},
16};
17
18use crate::{
19    Error, FlowInfo, FlowStage, LatencySample, LoadedLatencyReport, LossReport, LossSample,
20    MeasurementPathReport, NetBenchConfig, NetBenchEvent, NetBenchFlow, NetBenchProbeConfig,
21    NetBenchProbeReport, NetBenchReceiveStream, NetBenchReport, NetBenchSendStream,
22    NetBenchSession, NetBenchTelemetry, PROTOCOL_VERSION, PathKind, Result, SCHEMA_VERSION,
23    ThroughputDirection, ThroughputReport, ThroughputSample, TransportReport,
24    config::{LOADED_LATENCY_INTERVAL, MAX_PROBE_RATE_PER_SECOND, THROUGHPUT_SAMPLE_INTERVAL},
25    statistics::latency_report,
26    wire::{
27        Capabilities, ControlMessage, PROBE_MAGIC, Probe, ProbeKind, read_control, write_control,
28    },
29};
30
31const THROUGHPUT_CLEANUP_TIMEOUT: Duration = Duration::from_secs(5);
32type SendStream = Box<dyn NetBenchSendStream>;
33type RecvStream = Box<dyn NetBenchReceiveStream>;
34
35#[derive(Clone)]
36struct Connection(Arc<dyn NetBenchSession>);
37
38impl Connection {
39    fn remote_id(&self) -> String {
40        self.0.remote_peer_id()
41    }
42
43    fn telemetry(&self) -> NetBenchTelemetry {
44        self.0.telemetry()
45    }
46
47    fn max_datagram_size(&self) -> Option<usize> {
48        self.0.max_datagram_size()
49    }
50
51    async fn open_uni(&self) -> Result<SendStream> {
52        Ok(self.0.open_bi().await?.into_split().0)
53    }
54
55    async fn accept_uni(&self) -> Result<RecvStream> {
56        Ok(self.0.accept_bi().await?.into_split().1)
57    }
58
59    async fn send_datagram_wait(&self, bytes: Vec<u8>) -> Result<()> {
60        self.0.send_datagram(bytes).await
61    }
62
63    async fn read_datagram(&self) -> Result<Vec<u8>> {
64        self.0.read_datagram().await
65    }
66}
67
68/// Starts benchmark business flows on a caller-owned authenticated session.
69#[derive(Clone)]
70pub struct NetBenchInitiator;
71
72impl std::fmt::Debug for NetBenchInitiator {
73    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        formatter.debug_struct("NetBenchInitiator").finish()
75    }
76}
77
78impl Default for NetBenchInitiator {
79    fn default() -> Self {
80        Self::new()
81    }
82}
83
84impl NetBenchInitiator {
85    /// Creates an initiator which never dials or owns an endpoint/connection lifecycle.
86    #[must_use]
87    pub fn new() -> Self {
88        Self
89    }
90
91    /// Runs a benchmark and returns only its final result.
92    ///
93    /// # Errors
94    ///
95    /// Returns a negotiation, measurement, timeout, cancellation, or peer error.
96    pub async fn run(&self, flow: NetBenchFlow, config: NetBenchConfig) -> Result<NetBenchReport> {
97        self.start(flow, config).await?.result().await
98    }
99
100    /// Runs only the low-bandwidth latency and loss probes.
101    ///
102    /// This does not open upload or download throughput streams and remains available when the
103    /// peer denies bandwidth-saturating measurements.
104    ///
105    /// # Errors
106    ///
107    /// Returns a connection, negotiation, measurement, or peer error.
108    pub async fn run_probes(
109        &self,
110        flow: NetBenchFlow,
111        config: NetBenchProbeConfig,
112    ) -> Result<NetBenchProbeReport> {
113        let (event_tx, _event_rx) = mpsc::channel(128);
114        let overall_timeout = config.overall_timeout;
115        tokio::time::timeout(
116            overall_timeout,
117            run_probe_benchmark(flow, config, &EventSink(event_tx)),
118        )
119        .await
120        .map_err(|_| Error::Timeout {
121            stage: "probe benchmark",
122        })?
123    }
124
125    /// Starts a benchmark and returns an event stream plus final result handle.
126    ///
127    /// # Errors
128    ///
129    /// This method currently only fails if the benchmark task cannot be initialized.
130    #[allow(
131        clippy::unused_async,
132        reason = "keeps the documented start(...).await API"
133    )]
134    pub async fn start(&self, flow: NetBenchFlow, config: NetBenchConfig) -> Result<NetBenchTest> {
135        let (event_tx, event_rx) = mpsc::channel(128);
136        let events = EventSink(event_tx);
137        let overall_timeout = config.overall_timeout;
138        let stage = BenchmarkStage::default();
139        let task_stage = stage.clone();
140        let task = tokio::spawn(async move {
141            let result = if let Ok(result) = tokio::time::timeout(
142                overall_timeout,
143                run_benchmark(flow, config, &events, &task_stage),
144            )
145            .await
146            {
147                result
148            } else {
149                Err(Error::Timeout {
150                    stage: task_stage.current(),
151                })
152            };
153            if let Ok(report) = &result {
154                events.send(NetBenchEvent::Finished(report.clone()));
155            }
156            result
157        });
158        let abort = task.abort_handle();
159
160        Ok(NetBenchTest {
161            events: event_rx,
162            task: Some(task),
163            abort,
164        })
165    }
166}
167
168/// A running benchmark.
169#[derive(Debug)]
170pub struct NetBenchTest {
171    events: mpsc::Receiver<NetBenchEvent>,
172    task: Option<JoinHandle<Result<NetBenchReport>>>,
173    abort: AbortHandle,
174}
175
176impl NetBenchTest {
177    /// Waits for the next progress event.
178    pub async fn next(&mut self) -> Option<NetBenchEvent> {
179        self.events.recv().await
180    }
181
182    /// Cancels the benchmark task without closing the caller-owned connection.
183    ///
184    /// Call [`Self::abort_and_wait`] when the caller must prove that the internal task has exited.
185    pub fn cancel(&self) {
186        self.abort.abort();
187    }
188
189    /// Returns whether the internal benchmark task has finished.
190    #[must_use]
191    pub fn is_finished(&self) -> bool {
192        self.task.as_ref().is_none_or(JoinHandle::is_finished)
193    }
194
195    /// Cancels the benchmark and waits until its internal task has exited.
196    ///
197    /// # Errors
198    ///
199    /// Returns an error if the task panicked while cancellation was racing with completion.
200    pub async fn abort_and_wait(mut self) -> Result<()> {
201        self.cancel();
202        let Some(task) = self.task.take() else {
203            return Ok(());
204        };
205        match task.await {
206            Ok(_) => Ok(()),
207            Err(error) if error.is_cancelled() => Ok(()),
208            Err(error) => Err(Error::Protocol(format!(
209                "benchmark task failed while being cancelled: {error}"
210            ))),
211        }
212    }
213
214    /// Waits for the final report.
215    ///
216    /// # Errors
217    ///
218    /// Returns the terminal error from the benchmark task.
219    pub async fn result(mut self) -> Result<NetBenchReport> {
220        let task = self
221            .task
222            .take()
223            .ok_or_else(|| Error::Protocol("benchmark result was already consumed".to_owned()))?;
224        match task.await {
225            Ok(result) => result,
226            Err(error) if error.is_cancelled() => Err(Error::Cancelled),
227            Err(error) => Err(Error::Protocol(format!(
228                "benchmark task ended without a result: {error}"
229            ))),
230        }
231    }
232}
233
234impl Drop for NetBenchTest {
235    fn drop(&mut self) {
236        self.cancel();
237    }
238}
239
240#[derive(Clone, Debug)]
241struct BenchmarkStage(Arc<Mutex<&'static str>>);
242
243impl Default for BenchmarkStage {
244    fn default() -> Self {
245        Self(Arc::new(Mutex::new("flow setup")))
246    }
247}
248
249impl BenchmarkStage {
250    fn set(&self, stage: &'static str) {
251        *lock_unpoisoned(&self.0) = stage;
252    }
253
254    fn current(&self) -> &'static str {
255        *lock_unpoisoned(&self.0)
256    }
257}
258
259fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
260    mutex
261        .lock()
262        .unwrap_or_else(std::sync::PoisonError::into_inner)
263}
264
265#[derive(Clone)]
266struct EventSink(mpsc::Sender<NetBenchEvent>);
267
268impl EventSink {
269    fn send(&self, event: NetBenchEvent) {
270        let _ = self.0.try_send(event);
271    }
272}
273
274#[allow(
275    clippy::too_many_lines,
276    reason = "linear ordering documents the benchmark state machine"
277)]
278async fn run_benchmark(
279    flow: NetBenchFlow,
280    config: NetBenchConfig,
281    events: &EventSink,
282    stage: &BenchmarkStage,
283) -> Result<NetBenchReport> {
284    let (session, mut control_send, mut control_recv) = flow.into_parts();
285    let connection = Connection(session);
286    let total_started = Instant::now();
287    let connect_started = Instant::now();
288    let peer_id = connection.remote_id();
289
290    stage.set("protocol negotiation");
291    events.send(NetBenchEvent::FlowStage(FlowStage::ControlStreamReady));
292    write_control(
293        &mut control_send,
294        &ControlMessage::ClientHello {
295            protocol_versions: vec![PROTOCOL_VERSION],
296            capabilities: Capabilities {
297                datagram_probes: true,
298                loaded_latency: true,
299                path_stats: true,
300            },
301        },
302    )
303    .await?;
304    events.send(NetBenchEvent::FlowStage(FlowStage::ClientHelloSent));
305
306    let hello = read_control(&mut control_recv).await?;
307    let first_control_message_time = connect_started.elapsed();
308    let ControlMessage::ServerHello {
309        selected_version,
310        limits,
311        capabilities,
312    } = hello
313    else {
314        return Err(peer_or_protocol_error(hello));
315    };
316    if selected_version != PROTOCOL_VERSION {
317        return Err(Error::UnsupportedProtocolVersion);
318    }
319    events.send(NetBenchEvent::FlowStage(FlowStage::ServerHelloReceived {
320        version: selected_version,
321    }));
322    if !capabilities.datagram_probes {
323        return finish_with_error(
324            &mut control_send,
325            &mut control_recv,
326            Error::Protocol("peer does not support Datagram probes".to_owned()),
327        )
328        .await;
329    }
330    if let Err(error) = validate_against_server_limits(&config, limits) {
331        return finish_with_error(&mut control_send, &mut control_recv, error).await;
332    }
333
334    stage.set("path stabilization");
335    let path_at_flow_start = selected_path_kind(&connection);
336    let (initial_path, stabilization_time_to_direct) = stabilize_path(
337        &connection,
338        config.path_stabilization_timeout,
339        connect_started,
340        events,
341    )
342    .await;
343    events.send(NetBenchEvent::Ready(FlowInfo {
344        negotiation_time: first_control_message_time,
345        path: initial_path,
346    }));
347    let (path_monitor, path_monitor_state) =
348        spawn_path_monitor(connection.clone(), events.clone(), connect_started);
349
350    let before = TransportSnapshot::capture(&connection);
351    let mut test_id = 1_u64;
352
353    stage.set("idle latency");
354    events.send(NetBenchEvent::LatencyStarted);
355    write_control(
356        &mut control_send,
357        &ControlMessage::StartLatency {
358            test_id,
359            duration_ms: duration_ms(config.latency_duration),
360            interval_ms: duration_ms(config.latency_interval),
361        },
362    )
363    .await?;
364    let idle = measure_probes(
365        connection.clone(),
366        test_id,
367        config.latency_duration,
368        config.latency_interval,
369        config.probe_timeout,
370        ProbeEventMode::Latency(events.clone()),
371    )
372    .await?;
373    write_control(&mut control_send, &ControlMessage::StopTest { test_id }).await?;
374    let idle_latency = latency_report(&idle.rtts);
375    test_id += 1;
376
377    stage.set("loss probes");
378    events.send(NetBenchEvent::LossStarted);
379    write_control(
380        &mut control_send,
381        &ControlMessage::StartLoss {
382            test_id,
383            duration_ms: duration_ms(config.loss_duration),
384            rate_per_second: config.loss_rate_per_second,
385            timeout_ms: duration_ms(config.probe_timeout),
386        },
387    )
388    .await?;
389    let loss_interval =
390        Duration::from_secs_f64(1.0 / f64::from(config.loss_rate_per_second.max(1)));
391    let loss_samples = measure_probes(
392        connection.clone(),
393        test_id,
394        config.loss_duration,
395        loss_interval,
396        config.probe_timeout,
397        ProbeEventMode::Loss(events.clone()),
398    )
399    .await?;
400    write_control(&mut control_send, &ControlMessage::StopTest { test_id }).await?;
401    let loss = loss_samples.loss_report();
402    test_id += 1;
403
404    if !config.download_warmup.is_zero() {
405        stage.set("download warm-up");
406        events.send(NetBenchEvent::DownloadWarmupStarted);
407        let _ = download_phase(
408            &connection,
409            &mut control_send,
410            &mut control_recv,
411            test_id,
412            config.download_warmup,
413            config.parallel_streams,
414            config.chunk_size,
415            None,
416        )
417        .await?;
418        test_id += 1;
419    }
420
421    stage.set("download throughput");
422    events.send(NetBenchEvent::DownloadStarted);
423    let download_future = download_phase(
424        &connection,
425        &mut control_send,
426        &mut control_recv,
427        test_id,
428        config.download_duration,
429        config.parallel_streams,
430        config.chunk_size,
431        Some(events.clone()),
432    );
433    let download_probe_future = measure_probes(
434        connection.clone(),
435        test_id,
436        config.download_duration,
437        LOADED_LATENCY_INTERVAL,
438        config.probe_timeout,
439        ProbeEventMode::Latency(events.clone()),
440    );
441    let (download, download_loaded) = tokio::try_join!(download_future, download_probe_future)?;
442    test_id += 1;
443
444    if !config.upload_warmup.is_zero() {
445        stage.set("upload warm-up");
446        events.send(NetBenchEvent::UploadWarmupStarted);
447        if !matches!(
448            selected_path_kind(&connection),
449            PathKind::Direct | PathKind::DirectIpv4 | PathKind::DirectIpv6
450        ) {
451            let _ = stabilize_path(
452                &connection,
453                config.path_stabilization_timeout,
454                connect_started,
455                events,
456            )
457            .await;
458        }
459        let _ = upload_phase(
460            &connection,
461            &mut control_send,
462            &mut control_recv,
463            test_id,
464            config.upload_warmup,
465            config.parallel_streams,
466            config.chunk_size,
467            None,
468        )
469        .await?;
470        test_id += 1;
471    }
472
473    stage.set("upload throughput");
474    events.send(NetBenchEvent::UploadStarted);
475    let upload_future = upload_phase(
476        &connection,
477        &mut control_send,
478        &mut control_recv,
479        test_id,
480        config.upload_duration,
481        config.parallel_streams,
482        config.chunk_size,
483        Some(events.clone()),
484    );
485    let upload_probe_future = measure_probes(
486        connection.clone(),
487        test_id,
488        config.upload_duration,
489        LOADED_LATENCY_INTERVAL,
490        config.probe_timeout,
491        ProbeEventMode::Latency(events.clone()),
492    );
493    let (upload, upload_loaded) = tokio::try_join!(upload_future, upload_probe_future)?;
494
495    stage.set("finalization");
496    let after = TransportSnapshot::capture(&connection);
497    path_monitor.abort_and_wait().await;
498    let final_path = selected_path_kind(&connection);
499    let path_monitor_state = lock_unpoisoned(&path_monitor_state).clone();
500    let time_to_direct = if path_at_flow_start == PathKind::Relay {
501        stabilization_time_to_direct.or(path_monitor_state.time_to_direct)
502    } else {
503        None
504    };
505    let download_p50 = latency_report(&download_loaded.rtts).p50;
506    let upload_p50 = latency_report(&upload_loaded.rtts).p50;
507    let loaded_latency = LoadedLatencyReport {
508        idle_p50: idle_latency.p50,
509        download_p50,
510        download_increase: download_p50.saturating_sub(idle_latency.p50),
511        upload_p50,
512        upload_increase: upload_p50.saturating_sub(idle_latency.p50),
513    };
514
515    finish_flow(&mut control_send, &mut control_recv).await?;
516
517    Ok(NetBenchReport {
518        schema_version: SCHEMA_VERSION,
519        protocol_version: selected_version,
520        peer_id,
521        total_duration: total_started.elapsed(),
522        path: MeasurementPathReport {
523            negotiation_time: first_control_message_time,
524            initial_path,
525            final_path,
526            path_changed: path_at_flow_start != initial_path
527                || initial_path != final_path
528                || path_monitor_state.path_changed,
529            became_direct: time_to_direct.is_some(),
530            time_to_direct,
531        },
532        idle_latency,
533        loss,
534        download,
535        upload,
536        loaded_latency,
537        transport: after.delta(before),
538    })
539}
540
541#[allow(
542    clippy::too_many_lines,
543    reason = "mirrors the full benchmark negotiation while omitting all throughput phases"
544)]
545async fn run_probe_benchmark(
546    flow: NetBenchFlow,
547    config: NetBenchProbeConfig,
548    events: &EventSink,
549) -> Result<NetBenchProbeReport> {
550    let (session, mut control_send, mut control_recv) = flow.into_parts();
551    let connection = Connection(session);
552    let total_started = Instant::now();
553    let connect_started = Instant::now();
554    let peer_id = connection.remote_id();
555
556    events.send(NetBenchEvent::FlowStage(FlowStage::ControlStreamReady));
557    write_control(
558        &mut control_send,
559        &ControlMessage::ClientHello {
560            protocol_versions: vec![PROTOCOL_VERSION],
561            capabilities: Capabilities {
562                datagram_probes: true,
563                loaded_latency: false,
564                path_stats: true,
565            },
566        },
567    )
568    .await?;
569    events.send(NetBenchEvent::FlowStage(FlowStage::ClientHelloSent));
570
571    let hello = read_control(&mut control_recv).await?;
572    let first_control_message_time = connect_started.elapsed();
573    let ControlMessage::ServerHello {
574        selected_version,
575        limits,
576        capabilities,
577    } = hello
578    else {
579        return Err(peer_or_protocol_error(hello));
580    };
581    if selected_version != PROTOCOL_VERSION {
582        return Err(Error::UnsupportedProtocolVersion);
583    }
584    events.send(NetBenchEvent::FlowStage(FlowStage::ServerHelloReceived {
585        version: selected_version,
586    }));
587    if !capabilities.datagram_probes {
588        return finish_with_error(
589            &mut control_send,
590            &mut control_recv,
591            Error::Protocol("peer does not support Datagram probes".to_owned()),
592        )
593        .await;
594    }
595    if let Err(error) = validate_probes_against_server_limits(&config, limits) {
596        return finish_with_error(&mut control_send, &mut control_recv, error).await;
597    }
598
599    let path_at_flow_start = selected_path_kind(&connection);
600    let (initial_path, stabilization_time_to_direct) = stabilize_path(
601        &connection,
602        config.path_stabilization_timeout,
603        connect_started,
604        events,
605    )
606    .await;
607    events.send(NetBenchEvent::Ready(FlowInfo {
608        negotiation_time: first_control_message_time,
609        path: initial_path,
610    }));
611    let (path_monitor, path_monitor_state) =
612        spawn_path_monitor(connection.clone(), events.clone(), connect_started);
613    let before = TransportSnapshot::capture(&connection);
614
615    let mut test_id = 1_u64;
616    events.send(NetBenchEvent::LatencyStarted);
617    write_control(
618        &mut control_send,
619        &ControlMessage::StartLatency {
620            test_id,
621            duration_ms: duration_ms(config.latency_duration),
622            interval_ms: duration_ms(config.latency_interval),
623        },
624    )
625    .await?;
626    let idle = measure_probes(
627        connection.clone(),
628        test_id,
629        config.latency_duration,
630        config.latency_interval,
631        config.probe_timeout,
632        ProbeEventMode::Latency(events.clone()),
633    )
634    .await?;
635    write_control(&mut control_send, &ControlMessage::StopTest { test_id }).await?;
636    let idle_latency = latency_report(&idle.rtts);
637    test_id += 1;
638
639    events.send(NetBenchEvent::LossStarted);
640    write_control(
641        &mut control_send,
642        &ControlMessage::StartLoss {
643            test_id,
644            duration_ms: duration_ms(config.loss_duration),
645            rate_per_second: config.loss_rate_per_second,
646            timeout_ms: duration_ms(config.probe_timeout),
647        },
648    )
649    .await?;
650    let loss_interval =
651        Duration::from_secs_f64(1.0 / f64::from(config.loss_rate_per_second.max(1)));
652    let loss_samples = measure_probes(
653        connection.clone(),
654        test_id,
655        config.loss_duration,
656        loss_interval,
657        config.probe_timeout,
658        ProbeEventMode::Loss(events.clone()),
659    )
660    .await?;
661    write_control(&mut control_send, &ControlMessage::StopTest { test_id }).await?;
662    let loss = loss_samples.loss_report();
663
664    let after = TransportSnapshot::capture(&connection);
665    path_monitor.abort_and_wait().await;
666    let final_path = selected_path_kind(&connection);
667    let path_monitor_state = lock_unpoisoned(&path_monitor_state).clone();
668    let time_to_direct = if path_at_flow_start == PathKind::Relay {
669        stabilization_time_to_direct.or(path_monitor_state.time_to_direct)
670    } else {
671        None
672    };
673    finish_flow(&mut control_send, &mut control_recv).await?;
674
675    Ok(NetBenchProbeReport {
676        schema_version: SCHEMA_VERSION,
677        protocol_version: selected_version,
678        peer_id,
679        total_duration: total_started.elapsed(),
680        path: MeasurementPathReport {
681            negotiation_time: first_control_message_time,
682            initial_path,
683            final_path,
684            path_changed: path_at_flow_start != initial_path
685                || initial_path != final_path
686                || path_monitor_state.path_changed,
687            became_direct: time_to_direct.is_some(),
688            time_to_direct,
689        },
690        idle_latency,
691        loss,
692        transport: after.delta(before),
693    })
694}
695
696fn validate_against_server_limits(
697    config: &NetBenchConfig,
698    limits: crate::wire::ServerLimits,
699) -> Result<()> {
700    validate_probe_parameters(
701        config.latency_duration,
702        config.latency_interval,
703        config.loss_duration,
704        config.loss_rate_per_second,
705    )?;
706    if limits.parallel_streams == 0 {
707        return Err(Error::ThroughputDeniedByPeer);
708    }
709    if config.parallel_streams == 0 {
710        return Err(Error::Protocol(
711            "parallel stream count must be greater than zero".to_owned(),
712        ));
713    }
714    if config.chunk_size == 0 {
715        return Err(Error::Protocol(
716            "throughput chunk size must be greater than zero".to_owned(),
717        ));
718    }
719    if duration_ms(config.download_duration) == 0 || duration_ms(config.upload_duration) == 0 {
720        return Err(Error::Protocol(
721            "throughput durations must be at least one millisecond".to_owned(),
722        ));
723    }
724    let maximum = Duration::from_millis(u64::from(limits.test_duration_ms));
725    for duration in [
726        config.latency_duration,
727        config.loss_duration,
728        config.download_duration,
729        config.upload_duration,
730        config.download_warmup,
731        config.upload_warmup,
732    ] {
733        if !duration.is_zero() && duration_ms(duration) == 0 {
734            return Err(Error::Protocol(
735                "phase durations must be zero or at least one millisecond".to_owned(),
736            ));
737        }
738        if duration > maximum {
739            return Err(Error::DurationLimitExceeded {
740                requested: duration,
741                maximum,
742            });
743        }
744    }
745    if config.parallel_streams > limits.parallel_streams {
746        return Err(Error::Protocol(format!(
747            "requested {} streams but server allows {}",
748            config.parallel_streams, limits.parallel_streams
749        )));
750    }
751    if config.chunk_size > limits.chunk_size {
752        return Err(Error::Protocol(format!(
753            "requested {} byte chunks but server allows {}",
754            config.chunk_size, limits.chunk_size
755        )));
756    }
757    Ok(())
758}
759
760fn validate_probes_against_server_limits(
761    config: &NetBenchProbeConfig,
762    limits: crate::wire::ServerLimits,
763) -> Result<()> {
764    validate_probe_parameters(
765        config.latency_duration,
766        config.latency_interval,
767        config.loss_duration,
768        config.loss_rate_per_second,
769    )?;
770    let maximum = Duration::from_millis(u64::from(limits.test_duration_ms));
771    for duration in [config.latency_duration, config.loss_duration] {
772        if duration > maximum {
773            return Err(Error::DurationLimitExceeded {
774                requested: duration,
775                maximum,
776            });
777        }
778    }
779    Ok(())
780}
781
782fn validate_probe_parameters(
783    latency_duration: Duration,
784    latency_interval: Duration,
785    loss_duration: Duration,
786    loss_rate_per_second: u32,
787) -> Result<()> {
788    if duration_ms(latency_duration) == 0 || duration_ms(loss_duration) == 0 {
789        return Err(Error::Protocol(
790            "probe durations must be at least one millisecond".to_owned(),
791        ));
792    }
793    if duration_ms(latency_interval) == 0 {
794        return Err(Error::Protocol(
795            "latency probe interval must be at least one millisecond".to_owned(),
796        ));
797    }
798    if !(1..=MAX_PROBE_RATE_PER_SECOND).contains(&loss_rate_per_second) {
799        return Err(Error::Protocol(format!(
800            "loss probe rate must be within 1..={MAX_PROBE_RATE_PER_SECOND} per second"
801        )));
802    }
803    Ok(())
804}
805
806async fn finish_with_error<T>(
807    send: &mut SendStream,
808    recv: &mut RecvStream,
809    error: Error,
810) -> Result<T> {
811    finish_flow(send, recv).await?;
812    Err(error)
813}
814
815fn peer_or_protocol_error(message: ControlMessage) -> Error {
816    match message {
817        ControlMessage::Error {
818            code,
819            message: peer_message,
820        } => Error::Peer {
821            code: code as u16,
822            message: peer_message,
823        },
824        other => Error::Protocol(format!("expected ServerHello, received {other:?}")),
825    }
826}
827
828async fn stabilize_path(
829    connection: &Connection,
830    timeout: Duration,
831    connect_started: Instant,
832    events: &EventSink,
833) -> (PathKind, Option<Duration>) {
834    let deadline = Instant::now() + timeout;
835    let mut previous = None;
836    loop {
837        let kind = selected_path_kind(connection);
838        if previous != Some(kind) {
839            events.send(NetBenchEvent::FlowStage(FlowStage::PathObserved {
840                path: kind,
841            }));
842            previous = Some(kind);
843        }
844        if matches!(
845            kind,
846            PathKind::Direct | PathKind::DirectIpv4 | PathKind::DirectIpv6
847        ) {
848            let elapsed = connect_started.elapsed();
849            events.send(NetBenchEvent::FlowStage(FlowStage::DirectPathSelected {
850                elapsed,
851            }));
852            return (kind, Some(elapsed));
853        }
854        if Instant::now() >= deadline {
855            if kind == PathKind::Relay {
856                events.send(NetBenchEvent::FlowStage(FlowStage::RelayRetained {
857                    waited: timeout,
858                }));
859            }
860            return (kind, None);
861        }
862        tokio::time::sleep(Duration::from_millis(100)).await;
863    }
864}
865
866#[derive(Clone)]
867enum ProbeEventMode {
868    Latency(EventSink),
869    Loss(EventSink),
870}
871
872struct ProbeResults {
873    sent: u64,
874    received: u64,
875    duplicated: u64,
876    reordered: u64,
877    rtts: Vec<Duration>,
878}
879
880impl ProbeResults {
881    #[allow(
882        clippy::cast_precision_loss,
883        reason = "ratios are intentionally reported as f64"
884    )]
885    fn loss_report(&self) -> LossReport {
886        let timed_out = self.sent.saturating_sub(self.received);
887        LossReport {
888            sent: self.sent,
889            received: self.received,
890            timed_out,
891            duplicated: self.duplicated,
892            reordered: self.reordered,
893            timeout_ratio: if self.sent == 0 {
894                0.0
895            } else {
896                timed_out as f64 / self.sent as f64
897            },
898        }
899    }
900}
901
902async fn measure_probes(
903    connection: Connection,
904    test_id: u64,
905    duration: Duration,
906    interval: Duration,
907    timeout: Duration,
908    mode: ProbeEventMode,
909) -> Result<ProbeResults> {
910    if connection.max_datagram_size().is_none() {
911        return Err(Error::Protocol(
912            "QUIC Datagram is unavailable on this connection".to_owned(),
913        ));
914    }
915
916    let started = Instant::now();
917    let send_deadline = started + duration;
918    let final_deadline = send_deadline + timeout;
919    let mut ticker = tokio::time::interval(interval.max(Duration::from_millis(1)));
920    ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
921    let mut sent_at = HashMap::<u64, Instant>::new();
922    let mut received = HashSet::<u64>::new();
923    let mut sequence = 0_u64;
924    let mut duplicated = 0_u64;
925    let mut reordered = 0_u64;
926    let mut highest_received = None::<u64>;
927    let mut rtts = Vec::new();
928
929    loop {
930        if Instant::now() >= final_deadline {
931            break;
932        }
933        tokio::select! {
934            _ = ticker.tick(), if Instant::now() < send_deadline => {
935                let probe = Probe {
936                    magic: PROBE_MAGIC,
937                    test_id,
938                    sequence,
939                    kind: ProbeKind::Request,
940                };
941                let payload = postcard::to_allocvec(&probe)?;
942                match tokio::time::timeout_at(
943                    send_deadline,
944                    connection.send_datagram_wait(payload),
945                )
946                .await
947                {
948                    Ok(result) => result.map_err(Error::network)?,
949                    Err(_) => break,
950                }
951                sent_at.insert(sequence, Instant::now());
952                sequence += 1;
953            }
954            datagram = connection.read_datagram() => {
955                let bytes = datagram.map_err(Error::network)?;
956                let Ok(probe) = postcard::from_bytes::<Probe>(&bytes) else {
957                    continue;
958                };
959                if probe.magic != PROBE_MAGIC || probe.test_id != test_id || probe.kind != ProbeKind::Response {
960                    continue;
961                }
962                let Some(sent) = sent_at.get(&probe.sequence) else {
963                    continue;
964                };
965                if !received.insert(probe.sequence) {
966                    duplicated += 1;
967                    continue;
968                }
969                if highest_received.is_some_and(|highest| probe.sequence < highest) {
970                    reordered += 1;
971                }
972                highest_received = Some(highest_received.map_or(probe.sequence, |value| value.max(probe.sequence)));
973                let rtt = sent.elapsed();
974                rtts.push(rtt);
975                match &mode {
976                    ProbeEventMode::Latency(events) => events.send(NetBenchEvent::LatencySample(
977                        LatencySample { sequence: probe.sequence, rtt }
978                    )),
979                    ProbeEventMode::Loss(events) => events.send(NetBenchEvent::LossSample(
980                        LossSample {
981                            sent: sequence,
982                            received: received.len() as u64,
983                            outstanding: sequence.saturating_sub(received.len() as u64),
984                        }
985                    )),
986                }
987            }
988            () = tokio::time::sleep_until(final_deadline) => break,
989        }
990
991        if Instant::now() >= send_deadline && received.len() == sent_at.len() {
992            break;
993        }
994    }
995
996    Ok(ProbeResults {
997        sent: sequence,
998        received: received.len() as u64,
999        duplicated,
1000        reordered,
1001        rtts,
1002    })
1003}
1004
1005#[allow(clippy::too_many_arguments)]
1006async fn download_phase(
1007    connection: &Connection,
1008    control_send: &mut SendStream,
1009    control_recv: &mut RecvStream,
1010    test_id: u64,
1011    duration: Duration,
1012    streams: u16,
1013    chunk_size: u32,
1014    events: Option<EventSink>,
1015) -> Result<ThroughputReport> {
1016    write_control(
1017        control_send,
1018        &ControlMessage::StartDownload {
1019            test_id,
1020            duration_ms: duration_ms(duration),
1021            streams,
1022            chunk_size,
1023        },
1024    )
1025    .await?;
1026
1027    let total = Arc::new(AtomicU64::new(0));
1028    let mut download_streams = Vec::with_capacity(usize::from(streams));
1029    for _ in 0..streams {
1030        let mut stream = connection.accept_uni().await?;
1031        let mut magic = [0_u8; 1];
1032        stream.read_exact(&mut magic).await?;
1033        if magic[0] != 0x44 {
1034            return Err(Error::Protocol(
1035                "download stream has an invalid pre-measurement header".to_owned(),
1036            ));
1037        }
1038        download_streams.push(stream);
1039    }
1040    expect_ready(control_recv, test_id).await?;
1041    write_control(control_send, &ControlMessage::TestReady { test_id }).await?;
1042
1043    let started = Instant::now();
1044    let deadline = started + duration;
1045    let sample_task = spawn_sampler(
1046        Arc::clone(&total),
1047        started,
1048        ThroughputDirection::Download,
1049        events,
1050    );
1051    let mut tasks = JoinSet::new();
1052    for mut stream in download_streams {
1053        let total = Arc::clone(&total);
1054        tasks.spawn(async move {
1055            let mut buffer = vec![0_u8; 64 * 1024];
1056            while Instant::now() < deadline {
1057                match tokio::time::timeout_at(deadline, stream.read(&mut buffer)).await {
1058                    Ok(Ok(read)) if read > 0 && Instant::now() <= deadline => {
1059                        total.fetch_add(read as u64, Ordering::Relaxed);
1060                    }
1061                    Ok(Ok(_)) | Err(_) => break,
1062                    Ok(Err(error)) => return Err(error),
1063                }
1064            }
1065            stream.cancel();
1066            Result::<()>::Ok(())
1067        });
1068    }
1069    while let Some(result) = tasks.join_next().await {
1070        result.map_err(Error::network)??;
1071    }
1072    tokio::time::sleep_until(deadline).await;
1073    sample_task.abort_and_wait().await;
1074    let bytes = total.load(Ordering::Relaxed);
1075    tracing::debug!(
1076        test_id,
1077        direction = "download",
1078        cleanup_timeout = ?THROUGHPUT_CLEANUP_TIMEOUT,
1079        "waiting for throughput completion on the prioritized control stream"
1080    );
1081    tokio::time::timeout(
1082        THROUGHPUT_CLEANUP_TIMEOUT,
1083        expect_finished(control_recv, test_id),
1084    )
1085    .await
1086    .map_err(|_| Error::Timeout {
1087        stage: "download cleanup",
1088    })??;
1089    Ok(throughput_report(bytes, duration, streams))
1090}
1091
1092#[allow(
1093    clippy::cast_precision_loss,
1094    clippy::too_many_arguments,
1095    clippy::too_many_lines,
1096    reason = "receiver-confirmed live samples keep the upload state machine explicit"
1097)]
1098async fn upload_phase(
1099    connection: &Connection,
1100    control_send: &mut SendStream,
1101    control_recv: &mut RecvStream,
1102    test_id: u64,
1103    duration: Duration,
1104    streams: u16,
1105    chunk_size: u32,
1106    events: Option<EventSink>,
1107) -> Result<ThroughputReport> {
1108    write_control(
1109        control_send,
1110        &ControlMessage::StartUpload {
1111            test_id,
1112            duration_ms: duration_ms(duration),
1113            streams,
1114            chunk_size,
1115        },
1116    )
1117    .await?;
1118
1119    let mut upload_streams = Vec::with_capacity(usize::from(streams));
1120    for _ in 0..streams {
1121        let mut stream = connection.open_uni().await?;
1122        stream.write_all(&[0x55]).await?;
1123        upload_streams.push(stream);
1124    }
1125    expect_ready(control_recv, test_id).await?;
1126
1127    let started = Instant::now();
1128    let deadline = started + duration;
1129    let chunk = Arc::new(vec![
1130        0x5A;
1131        usize::try_from(chunk_size).map_err(Error::network)?
1132    ]);
1133    let mut tasks = JoinSet::new();
1134    for mut stream in upload_streams {
1135        let chunk = Arc::clone(&chunk);
1136        tasks.spawn(async move {
1137            while Instant::now() < deadline {
1138                match tokio::time::timeout_at(deadline, stream.write(&chunk)).await {
1139                    Ok(Ok(_)) => {}
1140                    Ok(Err(Error::FlowStopped)) | Err(_) => break,
1141                    Ok(Err(error)) => return Err(error),
1142                }
1143            }
1144            stream.cancel();
1145            Result::<()>::Ok(())
1146        });
1147    }
1148    let mut progress = VecDeque::<(Duration, u64)>::from([(Duration::ZERO, 0)]);
1149    let cleanup_deadline = deadline + THROUGHPUT_CLEANUP_TIMEOUT;
1150    tracing::debug!(
1151        test_id,
1152        direction = "upload",
1153        cleanup_timeout = ?THROUGHPUT_CLEANUP_TIMEOUT,
1154        "waiting for throughput completion on the prioritized control stream"
1155    );
1156    let received_bytes = loop {
1157        let message = tokio::time::timeout_at(cleanup_deadline, read_control(control_recv))
1158            .await
1159            .map_err(|_| Error::Timeout {
1160                stage: "upload cleanup",
1161            })??;
1162        match message {
1163            ControlMessage::ThroughputProgress {
1164                test_id: received_test_id,
1165                received_bytes,
1166                duration_ns,
1167            } if received_test_id == test_id => {
1168                let receiver_elapsed = Duration::from_nanos(duration_ns).min(duration);
1169                progress.push_back((receiver_elapsed, received_bytes));
1170                while progress.len() > 5 {
1171                    progress.pop_front();
1172                }
1173                let (old_elapsed, old_bytes) = progress.front().copied().unwrap_or_default();
1174                let sample_duration = receiver_elapsed.saturating_sub(old_elapsed);
1175                let interval_bps = if sample_duration.is_zero() {
1176                    0.0
1177                } else {
1178                    received_bytes.saturating_sub(old_bytes) as f64 * 8.0
1179                        / sample_duration.as_secs_f64()
1180                };
1181                if let Some(events) = &events {
1182                    events.send(NetBenchEvent::UploadSample(ThroughputSample {
1183                        direction: ThroughputDirection::Upload,
1184                        elapsed: receiver_elapsed,
1185                        received_bytes,
1186                        interval_bps,
1187                    }));
1188                }
1189            }
1190            ControlMessage::TestFinished {
1191                test_id: received_test_id,
1192                received_bytes,
1193                ..
1194            } if received_test_id == test_id => break received_bytes,
1195            ControlMessage::Error { code, message } => {
1196                return Err(Error::Peer {
1197                    code: code as u16,
1198                    message,
1199                });
1200            }
1201            message => {
1202                return Err(Error::Protocol(format!(
1203                    "expected upload progress or completion, received {message:?}"
1204                )));
1205            }
1206        }
1207    };
1208    tokio::time::timeout_at(cleanup_deadline, async {
1209        while let Some(result) = tasks.join_next().await {
1210            result.map_err(Error::network)??;
1211        }
1212        Result::<()>::Ok(())
1213    })
1214    .await
1215    .map_err(|_| Error::Timeout {
1216        stage: "upload stream cleanup",
1217    })??;
1218    Ok(throughput_report(received_bytes, duration, streams))
1219}
1220
1221async fn expect_ready(control_recv: &mut RecvStream, expected_test_id: u64) -> Result<()> {
1222    match read_control(control_recv).await? {
1223        ControlMessage::TestReady { test_id } if test_id == expected_test_id => Ok(()),
1224        ControlMessage::TestReady { test_id } => Err(Error::Protocol(format!(
1225            "received readiness for test {test_id}, expected {expected_test_id}"
1226        ))),
1227        ControlMessage::Error { code, message } => Err(Error::Peer {
1228            code: code as u16,
1229            message,
1230        }),
1231        message => Err(Error::Protocol(format!(
1232            "expected TestReady, received {message:?}"
1233        ))),
1234    }
1235}
1236
1237#[allow(
1238    clippy::cast_precision_loss,
1239    reason = "throughput samples are intentionally reported as f64"
1240)]
1241fn spawn_sampler(
1242    total: Arc<AtomicU64>,
1243    started: Instant,
1244    direction: ThroughputDirection,
1245    events: Option<EventSink>,
1246) -> AbortOnDropTask<()> {
1247    AbortOnDropTask::new(tokio::spawn(async move {
1248        let Some(events) = events else {
1249            return;
1250        };
1251        let mut samples = VecDeque::<(Instant, u64)>::from([(started, 0)]);
1252        let mut ticker = tokio::time::interval(THROUGHPUT_SAMPLE_INTERVAL);
1253        ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
1254        ticker.tick().await;
1255        loop {
1256            ticker.tick().await;
1257            let now = Instant::now();
1258            let bytes = total.load(Ordering::Relaxed);
1259            samples.push_back((now, bytes));
1260            while samples.len() > 5 {
1261                samples.pop_front();
1262            }
1263            let (previous_time, previous_bytes) = samples.front().copied().unwrap_or((now, bytes));
1264            let delta_bytes = bytes.saturating_sub(previous_bytes);
1265            let delta_seconds = now.duration_since(previous_time).as_secs_f64();
1266            let interval_bps = if delta_seconds > 0.0 {
1267                delta_bytes as f64 * 8.0 / delta_seconds
1268            } else {
1269                0.0
1270            };
1271            events.send(match direction {
1272                ThroughputDirection::Download => NetBenchEvent::DownloadSample(ThroughputSample {
1273                    direction,
1274                    elapsed: started.elapsed(),
1275                    received_bytes: bytes,
1276                    interval_bps,
1277                }),
1278                ThroughputDirection::Upload => NetBenchEvent::UploadSample(ThroughputSample {
1279                    direction,
1280                    elapsed: started.elapsed(),
1281                    received_bytes: bytes,
1282                    interval_bps,
1283                }),
1284            });
1285        }
1286    }))
1287}
1288
1289#[derive(Clone, Debug, Default)]
1290struct PathMonitorState {
1291    path_changed: bool,
1292    time_to_direct: Option<Duration>,
1293}
1294
1295fn spawn_path_monitor(
1296    connection: Connection,
1297    events: EventSink,
1298    flow_started: Instant,
1299) -> (AbortOnDropTask<()>, Arc<Mutex<PathMonitorState>>) {
1300    let state = Arc::new(Mutex::new(PathMonitorState::default()));
1301    let task_state = Arc::clone(&state);
1302    let task = AbortOnDropTask::new(tokio::spawn(async move {
1303        let mut previous = selected_path_kind(&connection);
1304        let mut ticker = tokio::time::interval(Duration::from_millis(250));
1305        loop {
1306            ticker.tick().await;
1307            let current = selected_path_kind(&connection);
1308            if current != previous {
1309                let mut state = lock_unpoisoned(&task_state);
1310                state.path_changed = true;
1311                if previous == PathKind::Relay
1312                    && matches!(
1313                        current,
1314                        PathKind::Direct | PathKind::DirectIpv4 | PathKind::DirectIpv6
1315                    )
1316                    && state.time_to_direct.is_none()
1317                {
1318                    state.time_to_direct = Some(flow_started.elapsed());
1319                }
1320                drop(state);
1321                events.send(NetBenchEvent::FlowStage(FlowStage::PathObserved {
1322                    path: current,
1323                }));
1324                previous = current;
1325            }
1326        }
1327    }));
1328    (task, state)
1329}
1330
1331#[derive(Debug)]
1332struct AbortOnDropTask<T>(Option<JoinHandle<T>>);
1333
1334impl<T> AbortOnDropTask<T> {
1335    fn new(task: JoinHandle<T>) -> Self {
1336        Self(Some(task))
1337    }
1338
1339    async fn abort_and_wait(mut self) {
1340        if let Some(task) = self.0.take() {
1341            task.abort();
1342            let _ = task.await;
1343        }
1344    }
1345}
1346
1347impl<T> Drop for AbortOnDropTask<T> {
1348    fn drop(&mut self) {
1349        if let Some(task) = &self.0 {
1350            task.abort();
1351        }
1352    }
1353}
1354
1355async fn expect_finished(recv: &mut RecvStream, expected_test_id: u64) -> Result<(u64, Duration)> {
1356    match read_control(recv).await? {
1357        ControlMessage::TestFinished {
1358            test_id,
1359            received_bytes,
1360            duration_ns,
1361        } if test_id == expected_test_id => Ok((received_bytes, Duration::from_nanos(duration_ns))),
1362        ControlMessage::Error { code, message } => Err(Error::Peer {
1363            code: code as u16,
1364            message,
1365        }),
1366        other => Err(Error::Protocol(format!(
1367            "expected TestFinished({expected_test_id}), received {other:?}"
1368        ))),
1369    }
1370}
1371
1372async fn finish_flow(send: &mut SendStream, recv: &mut RecvStream) -> Result<()> {
1373    write_control(send, &ControlMessage::FlowFinished).await?;
1374    match read_control(recv).await? {
1375        ControlMessage::FlowFinishedAck => {
1376            send.finish().map_err(Error::network)?;
1377            Ok(())
1378        }
1379        ControlMessage::Error { code, message } => Err(Error::Peer {
1380            code: code as u16,
1381            message,
1382        }),
1383        other => Err(Error::Protocol(format!(
1384            "expected FlowFinishedAck, received {other:?}"
1385        ))),
1386    }
1387}
1388
1389#[allow(
1390    clippy::cast_precision_loss,
1391    reason = "throughput is intentionally reported as f64"
1392)]
1393fn throughput_report(bytes: u64, duration: Duration, streams: u16) -> ThroughputReport {
1394    ThroughputReport {
1395        received_bytes: bytes,
1396        measurement_duration: duration,
1397        bits_per_second: if duration.is_zero() {
1398            0.0
1399        } else {
1400            bytes as f64 * 8.0 / duration.as_secs_f64()
1401        },
1402        streams,
1403    }
1404}
1405
1406fn selected_path_kind(connection: &Connection) -> PathKind {
1407    connection.telemetry().path
1408}
1409
1410#[derive(Clone, Copy, Default)]
1411struct TransportSnapshot {
1412    connection_lost_packets: u64,
1413    connection_lost_bytes: u64,
1414    udp_rx_datagrams: u64,
1415    udp_tx_datagrams: u64,
1416    congestion_events: u64,
1417    black_holes_detected: u64,
1418    current_mtu: u16,
1419    rtt: Duration,
1420}
1421
1422impl TransportSnapshot {
1423    fn capture(connection: &Connection) -> Self {
1424        let telemetry = connection.telemetry();
1425        Self {
1426            connection_lost_packets: telemetry.lost_packets,
1427            connection_lost_bytes: telemetry.lost_bytes,
1428            udp_rx_datagrams: telemetry.rx_datagrams,
1429            udp_tx_datagrams: telemetry.tx_datagrams,
1430            congestion_events: telemetry.congestion_events,
1431            black_holes_detected: telemetry.black_holes_detected,
1432            current_mtu: telemetry.current_mtu,
1433            rtt: telemetry.rtt,
1434        }
1435    }
1436
1437    fn delta(self, before: Self) -> TransportReport {
1438        TransportReport {
1439            lost_packets: self
1440                .connection_lost_packets
1441                .saturating_sub(before.connection_lost_packets),
1442            lost_bytes: self
1443                .connection_lost_bytes
1444                .saturating_sub(before.connection_lost_bytes),
1445            congestion_events: self
1446                .congestion_events
1447                .saturating_sub(before.congestion_events),
1448            udp_rx_datagrams: self
1449                .udp_rx_datagrams
1450                .saturating_sub(before.udp_rx_datagrams),
1451            udp_tx_datagrams: self
1452                .udp_tx_datagrams
1453                .saturating_sub(before.udp_tx_datagrams),
1454            current_mtu: self.current_mtu,
1455            black_holes_detected: self
1456                .black_holes_detected
1457                .saturating_sub(before.black_holes_detected),
1458            final_rtt: self.rtt,
1459        }
1460    }
1461}
1462
1463fn duration_ms(duration: Duration) -> u32 {
1464    u32::try_from(duration.as_millis()).unwrap_or(u32::MAX)
1465}
iroh_netbench/initiator.rs

iroh_netbench/
initiator.rs
