From 4770d0d4761c8a2f1d65820aebdee0aa88bd27fc Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Fri, 1 May 2026 22:10:57 +0700 Subject: [PATCH] feat(cli): integrate StatsCollector into main execution loop --- crates/irtt-cli/src/main.rs | 161 ++++++++++++++++++++++++++++++------ 1 file changed, 134 insertions(+), 27 deletions(-) diff --git a/crates/irtt-cli/src/main.rs b/crates/irtt-cli/src/main.rs index 7a2bfce..b781226 100644 --- a/crates/irtt-cli/src/main.rs +++ b/crates/irtt-cli/src/main.rs @@ -8,6 +8,8 @@ use std::{ use clap::Parser; use irtt_cli::{format_event, CliArgs}; use irtt_client::{Client, ClientEvent, ClientTimestamp, OpenOutcome, RecvBudget}; +#[cfg(feature = "stats")] +use irtt_stats::{StatsCollector, StatsConfig}; const RECV_BUDGET: RecvBudget = RecvBudget { max_packets: 16 }; // TODO: derive final drain from pending probe timeout behavior once the client @@ -29,35 +31,44 @@ fn main() -> ExitCode { fn run(args: CliArgs) -> Result<(), Box> { let mode = args.output; let mut stdout = io::LineWriter::new(io::stdout().lock()); + #[cfg(feature = "stats")] + let mut stats = StatsCollector::new(StatsConfig::finite()); + let mut output = EventOutput { + mode, + out: &mut stdout, + #[cfg(feature = "stats")] + stats: &mut stats, + }; let mut client = Client::connect(args.to_client_config())?; let open = client.open(ClientTimestamp::now())?; - print_event(open_event(&open), mode, &mut stdout)?; + output.print_event(open_event(&open))?; while let Some(deadline) = client.next_send_deadline() { let now = Instant::now(); if deadline <= now { let events = client.send_probe()?; - print_events(&events, mode, &mut stdout)?; + output.print_events(&events)?; } let events = client.recv_available(RECV_BUDGET)?; - print_events(&events, mode, &mut stdout)?; + output.print_events(&events)?; let events = client.poll_timeouts(ClientTimestamp::now())?; - print_events(&events, mode, &mut stdout)?; + output.print_events(&events)?; sleep_until_next_send(client.next_send_deadline()); } - drain_final_replies(&mut client, mode, &mut stdout)?; + drain_final_replies(&mut client, &mut output)?; let events = client.poll_timeouts(ClientTimestamp::now())?; - print_events(&events, mode, &mut stdout)?; + output.print_events(&events)?; let events = client.close(ClientTimestamp::now())?; - print_events(&events, mode, &mut stdout)?; - stdout.flush()?; + output.print_events(&events)?; + output.print_summary()?; + output.out.flush()?; Ok(()) } @@ -69,8 +80,7 @@ fn open_event(outcome: &OpenOutcome) -> &ClientEvent { fn drain_final_replies( client: &mut Client, - mode: irtt_cli::OutputMode, - out: &mut W, + output: &mut EventOutput<'_, W>, ) -> Result<(), Box> { let deadline = Instant::now() + FINAL_DRAIN; while Instant::now() < deadline { @@ -78,7 +88,7 @@ fn drain_final_replies( if events.is_empty() { thread::sleep(Duration::from_millis(5)); } else { - print_events(&events, mode, out)?; + output.print_events(&events)?; } } Ok(()) @@ -96,24 +106,121 @@ fn sleep_until_next_send(deadline: Option) { } } -fn print_events( - events: &[ClientEvent], +struct EventOutput<'a, W: Write> { mode: irtt_cli::OutputMode, - out: &mut W, -) -> io::Result<()> { - for event in events { - print_event(event, mode, out)?; - } - Ok(()) + out: &'a mut W, + #[cfg(feature = "stats")] + stats: &'a mut StatsCollector, } -fn print_event( - event: &ClientEvent, - mode: irtt_cli::OutputMode, - out: &mut W, -) -> io::Result<()> { - if let Some(line) = format_event(event, mode) { - writeln!(out, "{line}")?; +impl EventOutput<'_, W> { + fn print_events(&mut self, events: &[ClientEvent]) -> io::Result<()> { + for event in events { + self.print_event(event)?; + } + Ok(()) + } + + fn print_event(&mut self, event: &ClientEvent) -> io::Result<()> { + #[cfg(feature = "stats")] + self.stats.process(event); + + if let Some(line) = format_event(event, self.mode) { + writeln!(self.out, "{line}")?; + } + Ok(()) + } + + fn print_summary(&mut self) -> io::Result<()> { + if !self.mode.prints_summary() { + return Ok(()); + } + + #[cfg(feature = "stats")] + { + write!( + self.out, + "{}", + irtt_cli::summary::format_summary(&self.stats.summary()) + )?; + } + + Ok(()) + } +} + +#[cfg(all(test, feature = "stats"))] +mod tests { + use super::*; + use irtt_client::{PacketMeta, RttSample, SignedDuration}; + use std::{ + net::{IpAddr, Ipv4Addr, SocketAddr}, + time::UNIX_EPOCH, + }; + + fn test_timestamp(offset: Duration) -> ClientTimestamp { + ClientTimestamp { + wall: UNIX_EPOCH + offset, + mono: Instant::now() + offset, + } + } + + fn test_remote() -> SocketAddr { + SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 2112) + } + + #[test] + fn output_helper_streams_and_collects_events() { + let sent_at = test_timestamp(Duration::from_secs(1)); + let received_at = test_timestamp(Duration::from_secs(1) + Duration::from_micros(1200)); + let events = [ + ClientEvent::EchoSent { + seq: 1, + logical_seq: 1, + remote: test_remote(), + scheduled_at: sent_at.mono, + sent_at, + bytes: 64, + send_call: Duration::from_micros(10), + timer_error: Duration::ZERO, + }, + ClientEvent::EchoReply { + seq: 1, + logical_seq: 1, + remote: test_remote(), + sent_at, + received_at, + rtt: RttSample { + raw: Duration::from_micros(1200), + adjusted: None, + effective: Duration::from_micros(1200), + adjusted_signed: None, + effective_signed: SignedDuration { ns: 1_200_000 }, + }, + server_timing: None, + one_way: None, + received_stats: None, + bytes: 64, + packet_meta: PacketMeta::default(), + }, + ]; + + let mut stats = StatsCollector::new(StatsConfig::finite()); + let mut out = Vec::new(); + { + let mut output = EventOutput { + mode: irtt_cli::OutputMode::RttUs, + out: &mut out, + stats: &mut stats, + }; + output.print_events(&events).unwrap(); + output.print_summary().unwrap(); + } + + let rendered = String::from_utf8(out).unwrap(); + let summary = stats.summary(); + assert_eq!(rendered, "1200\n"); + assert_eq!(summary.packets.packets_sent, 1); + assert_eq!(summary.packets.unique_replies, 1); } - Ok(()) }