feat(cli): integrate StatsCollector into main execution loop

This commit is contained in:
Nils Andreas Svee committed 2026-05-01 22:11:47 +07:00
1 parent 815dab077d
commit a26e1a6d59
1 file changed
+134 -27
+134 -27
View File
@@ -8,6 +8,8 @@ use std::{
use clap::Parser; use clap::Parser;
use irtt_cli::{format_event, CliArgs}; use irtt_cli::{format_event, CliArgs};
use irtt_client::{Client, ClientEvent, ClientTimestamp, OpenOutcome, RecvBudget}; use irtt_client::{Client, ClientEvent, ClientTimestamp, OpenOutcome, RecvBudget};
#[cfg(feature = "stats")]
use irtt_stats::{StatsCollector, StatsConfig};
const RECV_BUDGET: RecvBudget = RecvBudget { max_packets: 16 }; const RECV_BUDGET: RecvBudget = RecvBudget { max_packets: 16 };
// TODO: derive final drain from pending probe timeout behavior once the client // 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<dyn std::error::Error>> { fn run(args: CliArgs) -> Result<(), Box<dyn std::error::Error>> {
let mode = args.output; let mode = args.output;
let mut stdout = io::LineWriter::new(io::stdout().lock()); 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 mut client = Client::connect(args.to_client_config())?;
let open = client.open(ClientTimestamp::now())?; 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() { while let Some(deadline) = client.next_send_deadline() {
let now = Instant::now(); let now = Instant::now();
if deadline <= now { if deadline <= now {
let events = client.send_probe()?; let events = client.send_probe()?;
print_events(&events, mode, &mut stdout)?; output.print_events(&events)?;
} }
let events = client.recv_available(RECV_BUDGET)?; let events = client.recv_available(RECV_BUDGET)?;
print_events(&events, mode, &mut stdout)?; output.print_events(&events)?;
let events = client.poll_timeouts(ClientTimestamp::now())?; 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()); 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())?; let events = client.poll_timeouts(ClientTimestamp::now())?;
print_events(&events, mode, &mut stdout)?; output.print_events(&events)?;
let events = client.close(ClientTimestamp::now())?; let events = client.close(ClientTimestamp::now())?;
print_events(&events, mode, &mut stdout)?; output.print_events(&events)?;
stdout.flush()?; output.print_summary()?;
output.out.flush()?;
Ok(()) Ok(())
} }
@@ -69,8 +80,7 @@ fn open_event(outcome: &OpenOutcome) -> &ClientEvent {
fn drain_final_replies<W: Write>( fn drain_final_replies<W: Write>(
client: &mut Client, client: &mut Client,
mode: irtt_cli::OutputMode, output: &mut EventOutput<'_, W>,
out: &mut W,
) -> Result<(), Box<dyn std::error::Error>> { ) -> Result<(), Box<dyn std::error::Error>> {
let deadline = Instant::now() + FINAL_DRAIN; let deadline = Instant::now() + FINAL_DRAIN;
while Instant::now() < deadline { while Instant::now() < deadline {
@@ -78,7 +88,7 @@ fn drain_final_replies<W: Write>(
if events.is_empty() { if events.is_empty() {
thread::sleep(Duration::from_millis(5)); thread::sleep(Duration::from_millis(5));
} else { } else {
print_events(&events, mode, out)?; output.print_events(&events)?;
} }
} }
Ok(()) Ok(())
@@ -96,24 +106,121 @@ fn sleep_until_next_send(deadline: Option<Instant>) {
} }
} }
fn print_events<W: Write>( struct EventOutput<'a, W: Write> {
events: &[ClientEvent],
mode: irtt_cli::OutputMode, mode: irtt_cli::OutputMode,
out: &mut W, out: &'a mut W,
) -> io::Result<()> { #[cfg(feature = "stats")]
for event in events { stats: &'a mut StatsCollector,
print_event(event, mode, out)?;
}
Ok(())
} }
fn print_event<W: Write>( impl<W: Write> EventOutput<'_, W> {
event: &ClientEvent, fn print_events(&mut self, events: &[ClientEvent]) -> io::Result<()> {
mode: irtt_cli::OutputMode, for event in events {
out: &mut W, self.print_event(event)?;
) -> io::Result<()> { }
if let Some(line) = format_event(event, mode) { Ok(())
writeln!(out, "{line}")?; }
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(())
} }