feat(cli): implement multi-applet binary dispatch and signal handling

This commit is contained in:
Nils Andreas Svee committed 2026-06-09 23:23:50 +07:00
1 parent d7cdcc5cee
commit 133ff71b72
4 files changed
+166 -608

No files matched your search

+61 -608
View File
@@ -1,57 +1,19 @@
use std::{
io::{self, Write},
env,
ffi::OsString,
process::ExitCode,
sync::{
atomic::{AtomicBool, Ordering},
Arc,
},
thread,
time::{Duration, Instant},
sync::{atomic::AtomicBool, Arc},
};
use clap::Parser;
use irtt_cli::{
format_event, format_human_event_with_options, CliArgs, HumanOutputOptions, OutputMode,
applet::{
default_applet_for_features, detect_applet_from_argv0, AppletAvailability, RequestedApplet,
},
signal::install_signal_handler,
};
use irtt_client::{Client, ClientEvent, OpenOutcome, RecvBudget};
#[cfg(all(test, feature = "stats"))]
use irtt_client::ClientTimestamp;
#[cfg(feature = "tui")]
mod tui;
#[cfg(feature = "tui")]
use tui::run_tui;
#[cfg(feature = "stats")]
use irtt_stats::{StatsCollector, StatsConfig};
const RECV_BUDGET: RecvBudget = RecvBudget { max_packets: 16 };
const MAX_FINAL_DRAIN: Duration = Duration::from_secs(30);
const IDLE_SLEEP: Duration = Duration::from_millis(5);
const MAX_SLEEP: Duration = Duration::from_millis(20);
#[cfg(feature = "stats")]
const FINITE_STATS_BYTES_PER_PROBE: u64 = 500;
#[cfg(feature = "stats")]
const MIB: u64 = 1024 * 1024;
#[cfg(feature = "stats")]
const GIB: u64 = 1024 * MIB;
#[cfg(feature = "stats")]
const FINITE_STATS_MEMORY_WARNING_BYTES: u64 = 128 * MIB;
#[cfg(feature = "stats")]
const FINITE_STATS_MEMORY_STRONG_WARNING_BYTES: u64 = 512 * MIB;
#[cfg(feature = "stats")]
const FINITE_STATS_MEMORY_VERY_STRONG_WARNING_BYTES: u64 = GIB;
fn main() -> ExitCode {
let args = CliArgs::parse();
let shutdown_requested = Arc::new(AtomicBool::new(false));
if let Err(err) = install_signal_handler(Arc::clone(&shutdown_requested)) {
eprintln!("irtt-rs: failed to install signal handler: {err}");
return ExitCode::FAILURE;
}
match run(args, shutdown_requested.as_ref()) {
match run_from_env() {
Ok(()) => ExitCode::SUCCESS,
Err(err) => {
eprintln!("irtt-rs: {err}");
@@ -60,580 +22,71 @@ fn main() -> ExitCode {
}
}
fn install_signal_handler(shutdown_requested: Arc<AtomicBool>) -> Result<(), ctrlc::Error> {
ctrlc::set_handler(move || {
shutdown_requested.store(true, Ordering::Relaxed);
})
}
fn run_from_env() -> Result<(), Box<dyn std::error::Error>> {
let argv: Vec<OsString> = env::args_os().collect();
let argv0 = argv
.first()
.and_then(|arg| arg.to_str())
.unwrap_or("irtt-rs");
let requested = match detect_applet_from_argv0(argv0) {
Some(applet) => applet,
None => default_requested_applet()?,
};
fn run(args: CliArgs, shutdown_requested: &AtomicBool) -> Result<(), Box<dyn std::error::Error>> {
match args.output {
OutputMode::Human | OutputMode::Machine | OutputMode::Simple | OutputMode::RttUs => {
run_stream(args, shutdown_requested)
}
#[cfg(feature = "tui")]
OutputMode::Tui => run_tui(args, shutdown_requested),
let shutdown_requested = Arc::new(AtomicBool::new(false));
install_signal_handler(Arc::clone(&shutdown_requested))
.map_err(|err| format!("failed to install signal handler: {err}"))?;
match requested {
RequestedApplet::Client => run_client_applet(argv, shutdown_requested.as_ref()),
RequestedApplet::Tui => run_tui_applet(argv, shutdown_requested.as_ref()),
RequestedApplet::Server => Err("server applet is not available in this build".into()),
}
}
fn run_stream(
args: CliArgs,
fn default_requested_applet() -> Result<RequestedApplet, Box<dyn std::error::Error>> {
match default_applet_for_features(cfg!(feature = "client"), cfg!(feature = "tui"), false) {
AppletAvailability::Client => Ok(RequestedApplet::Client),
AppletAvailability::Tui => Ok(RequestedApplet::Tui),
AppletAvailability::Server => Ok(RequestedApplet::Server),
AppletAvailability::None => Err("no runnable applet is available in this build".into()),
}
}
#[cfg(feature = "client")]
fn run_client_applet(
argv: Vec<OsString>,
shutdown_requested: &AtomicBool,
) -> Result<(), Box<dyn std::error::Error>> {
let mode = args.output;
let continuous = args.is_continuous();
#[cfg(feature = "stats")]
if let Some(warning) = finite_stats_memory_warning(&args) {
eprintln!("{warning}");
}
let mut stdout = io::LineWriter::new(io::stdout().lock());
#[cfg(feature = "stats")]
let mut stats = StatsCollector::new(stats_config(continuous));
let mut stream_output = StreamOutput {
mode,
human_options: HumanOutputOptions {
verbose: args.verbose,
},
print_final_summary: false,
show_running_only_summary_note: false,
out: &mut stdout,
#[cfg(feature = "stats")]
stats: &mut stats,
};
use clap::Parser;
if is_shutdown_requested(shutdown_requested) {
return Ok(());
}
let mut client = Client::connect(args.to_client_config())?;
if is_shutdown_requested(shutdown_requested) {
return Ok(());
}
let open = client.open()?;
stream_output.print_event(open_event(&open))?;
let mut interrupted = false;
while should_continue_run(continuous, client.next_send_deadline(), shutdown_requested) {
if should_send_probe(client.next_send_deadline(), shutdown_requested) {
let events = client.send_probe()?;
stream_output.print_events(&events)?;
}
let events = client.recv_available(RECV_BUDGET)?;
stream_output.print_events(&events)?;
let events = client.poll_timeouts()?;
stream_output.print_events(&events)?;
if is_shutdown_requested(shutdown_requested) {
interrupted = true;
break;
}
sleep_until_next_send(client.next_send_deadline());
}
interrupted |= is_shutdown_requested(shutdown_requested);
if interrupted {
eprintln!("interrupted, closing session...");
}
if should_drain_final(continuous, interrupted) {
drain_final_replies(&mut client, &mut stream_output)?;
}
let events = client.poll_timeouts()?;
stream_output.print_events(&events)?;
let events = client.close()?;
stream_output.print_events(&events)?;
let print_final_summary = should_print_final_summary(continuous, interrupted);
stream_output.print_final_summary = print_final_summary;
stream_output.show_running_only_summary_note = continuous && interrupted && print_final_summary;
stream_output.print_summary()?;
stream_output.out.flush()?;
Ok(())
let args = irtt_cli::cmd::client::ClientArgs::parse_from(argv);
irtt_cli::cmd::client::run_stream(args, shutdown_requested)
}
fn is_shutdown_requested(shutdown_requested: &AtomicBool) -> bool {
shutdown_requested.load(Ordering::Relaxed)
}
fn should_continue_run(
continuous: bool,
next_send_deadline: Option<Instant>,
shutdown_requested: &AtomicBool,
) -> bool {
!is_shutdown_requested(shutdown_requested) && (continuous || next_send_deadline.is_some())
}
fn should_send_probe(next_send_deadline: Option<Instant>, shutdown_requested: &AtomicBool) -> bool {
!is_shutdown_requested(shutdown_requested)
&& next_send_deadline.is_some_and(|deadline| deadline <= Instant::now())
}
fn should_drain_final(continuous: bool, interrupted: bool) -> bool {
!continuous || interrupted
}
fn should_print_final_summary(continuous: bool, interrupted: bool) -> bool {
!continuous || interrupted
}
fn open_event(outcome: &OpenOutcome) -> &ClientEvent {
match outcome {
OpenOutcome::Started { event, .. } | OpenOutcome::NoTestCompleted { event, .. } => event,
}
}
fn drain_final_replies<W: Write>(
client: &mut Client,
stream_output: &mut StreamOutput<'_, W>,
#[cfg(not(feature = "client"))]
fn run_client_applet(
_argv: Vec<OsString>,
_shutdown_requested: &AtomicBool,
) -> Result<(), Box<dyn std::error::Error>> {
let deadline = Instant::now() + final_drain_duration(client.probe_timeout());
loop {
let mut printed = false;
let events = client.recv_available(RECV_BUDGET)?;
printed |= !events.is_empty();
stream_output.print_events(&events)?;
let events = client.poll_timeouts()?;
printed |= !events.is_empty();
stream_output.print_events(&events)?;
if client.is_run_complete() || Instant::now() >= deadline {
break;
}
if !printed {
thread::sleep(IDLE_SLEEP);
}
}
Ok(())
Err("client applet is not available; rebuild with the client feature".into())
}
fn sleep_until_next_send(deadline: Option<Instant>) {
let sleep_for = match deadline {
Some(deadline) => deadline
.saturating_duration_since(Instant::now())
.min(MAX_SLEEP),
None => Duration::from_millis(1),
};
if !sleep_for.is_zero() {
thread::sleep(sleep_for);
}
#[cfg(feature = "tui")]
fn run_tui_applet(
argv: Vec<OsString>,
shutdown_requested: &AtomicBool,
) -> Result<(), Box<dyn std::error::Error>> {
use clap::Parser;
let args = irtt_cli::cmd::tui::TuiArgs::parse_from(argv);
irtt_cli::cmd::tui::run_tui(args, shutdown_requested)
}
struct StreamOutput<'a, W: Write> {
mode: irtt_cli::OutputMode,
human_options: HumanOutputOptions,
print_final_summary: bool,
show_running_only_summary_note: bool,
out: &'a mut W,
#[cfg(feature = "stats")]
stats: &'a mut StatsCollector,
}
impl<W: Write> StreamOutput<'_, 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")]
let stats_update = self.stats.process(event);
#[cfg(not(feature = "stats"))]
let stats_update = irtt_cli::HumanEventStats::default();
let line =
if self.mode == OutputMode::Human && !matches!(event, ClientEvent::EchoSent { .. }) {
#[cfg(feature = "stats")]
let human_stats = stats_update.into();
#[cfg(not(feature = "stats"))]
let human_stats = stats_update;
Some(format_human_event_with_options(
event,
Some(human_stats),
self.human_options,
))
} else {
format_event(event, self.mode)
};
if let Some(line) = line {
writeln!(self.out, "{line}")?;
}
Ok(())
}
fn print_summary(&mut self) -> io::Result<()> {
if !self.print_final_summary || !self.mode.prints_summary() {
return Ok(());
}
#[cfg(feature = "stats")]
{
write!(
self.out,
"{}",
irtt_cli::summary::format_summary_with_options(
&self.stats.snapshot(),
irtt_cli::summary::SummaryFormatOptions {
verbose: self.human_options.verbose,
show_running_only_note: self.show_running_only_summary_note,
},
)
)?;
}
Ok(())
}
}
#[cfg(feature = "stats")]
fn stats_config(continuous: bool) -> StatsConfig {
if continuous {
StatsConfig::continuous()
} else {
StatsConfig::finite()
}
}
#[cfg(feature = "stats")]
fn expected_probe_count(duration: Duration, interval: Duration) -> u64 {
let interval_nanos = interval.as_nanos();
if interval_nanos == 0 {
return u64::MAX;
}
let expected = duration
.as_nanos()
.saturating_add(interval_nanos.saturating_sub(1))
/ interval_nanos;
expected.min(u128::from(u64::MAX)) as u64
}
#[cfg(feature = "stats")]
fn estimate_finite_stats_memory_bytes(expected_probes: u64) -> u64 {
expected_probes.saturating_mul(FINITE_STATS_BYTES_PER_PROBE)
}
#[cfg(feature = "stats")]
fn finite_stats_memory_warning(args: &CliArgs) -> Option<String> {
if args.is_continuous() || args.duration.is_zero() {
return None;
}
let expected_probes = expected_probe_count(args.duration, args.interval);
let estimated_bytes = estimate_finite_stats_memory_bytes(expected_probes);
if estimated_bytes < FINITE_STATS_MEMORY_WARNING_BYTES {
return None;
}
let formatted = format_bytes_for_warning(estimated_bytes);
let guidance = if estimated_bytes >= FINITE_STATS_MEMORY_VERY_STRONG_WARNING_BYTES {
"this may be unsuitable on memory-constrained systems"
} else if estimated_bytes >= FINITE_STATS_MEMORY_STRONG_WARNING_BYTES {
"consider shortening the run, increasing the interval, or using continuous mode"
} else {
"use continuous mode for bounded-memory long-running tests"
};
Some(format!(
"irtt-rs: warning: finite exact statistics may retain about {formatted} for this run; {guidance}"
))
}
#[cfg(feature = "stats")]
fn format_bytes_for_warning(bytes: u64) -> String {
if bytes >= GIB {
format!("{} GiB", bytes.saturating_add(GIB / 2) / GIB)
} else {
format!("{} MiB", bytes.saturating_add(MIB / 2) / MIB)
}
}
fn final_drain_duration(probe_timeout: Duration) -> Duration {
probe_timeout.min(MAX_FINAL_DRAIN)
}
#[cfg(test)]
mod shutdown_tests {
use super::*;
#[test]
fn final_drain_uses_capped_probe_timeout() {
assert_eq!(
final_drain_duration(Duration::from_secs(4)),
Duration::from_secs(4)
);
assert_eq!(
final_drain_duration(Duration::from_secs(60)),
Duration::from_secs(30)
);
}
}
#[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)
}
fn reply_event(seq: u32, rtt_us: u64) -> ClientEvent {
let sent_at = test_timestamp(Duration::from_millis(seq as u64));
let received_at = ClientTimestamp {
wall: sent_at.wall + Duration::from_micros(rtt_us),
mono: sent_at.mono + Duration::from_micros(rtt_us),
};
ClientEvent::EchoReply {
seq,
remote: test_remote(),
sent_at,
received_at,
rtt: RttSample {
raw: Duration::from_micros(rtt_us),
adjusted: None,
effective: SignedDuration::from_nanos(i128::from(rtt_us) * 1_000),
},
server_timing: None,
one_way: None,
received_stats: None,
bytes: 64,
packet_meta: PacketMeta::default(),
}
}
fn cli_args(args: &[&str]) -> CliArgs {
let mut argv = vec!["irtt-rs"];
argv.extend_from_slice(args);
CliArgs::try_parse_from(argv).unwrap()
}
#[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,
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,
remote: test_remote(),
sent_at,
received_at,
rtt: RttSample {
raw: Duration::from_micros(1200),
adjusted: None,
effective: SignedDuration::from_nanos(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 stream_output = StreamOutput {
mode: irtt_cli::OutputMode::RttUs,
human_options: HumanOutputOptions::default(),
print_final_summary: true,
show_running_only_summary_note: false,
out: &mut out,
stats: &mut stats,
};
stream_output.print_events(&events).unwrap();
stream_output.print_summary().unwrap();
}
let rendered = String::from_utf8(out).unwrap();
let summary = stats.snapshot();
assert_eq!(rendered, "1200\n");
assert_eq!(summary.packets.packets_sent, 1);
assert_eq!(summary.packets.unique_replies, 1);
}
#[test]
fn continuous_summary_prints_when_enabled_and_suppresses_when_disabled() {
let mut stats = StatsCollector::new(StatsConfig::continuous());
let mut out = Vec::new();
{
let mut stream_output = StreamOutput {
mode: irtt_cli::OutputMode::Human,
human_options: HumanOutputOptions::default(),
print_final_summary: true,
show_running_only_summary_note: true,
out: &mut out,
stats: &mut stats,
};
stream_output.print_event(&reply_event(1, 1200)).unwrap();
stream_output.print_summary().unwrap();
}
let rendered = String::from_utf8(out).unwrap();
assert!(rendered.contains("irtt-rs summary"));
assert!(rendered.contains("medians unavailable"));
assert!(rendered.contains("continuous mode"));
assert!(rendered.contains("packets:"));
assert!(rendered.contains("received=1"));
let mut stats = StatsCollector::new(StatsConfig::continuous());
let mut out = Vec::new();
{
let mut stream_output = StreamOutput {
mode: irtt_cli::OutputMode::Human,
human_options: HumanOutputOptions::default(),
print_final_summary: false,
show_running_only_summary_note: true,
out: &mut out,
stats: &mut stats,
};
stream_output.print_summary().unwrap();
}
assert!(out.is_empty());
}
#[test]
fn human_output_prints_ipdv_from_stats_update() {
let events = [reply_event(0, 1200), reply_event(1, 1250)];
let mut stats = StatsCollector::new(StatsConfig::finite());
let mut out = Vec::new();
{
let mut stream_output = StreamOutput {
mode: irtt_cli::OutputMode::Human,
human_options: HumanOutputOptions::default(),
print_final_summary: false,
show_running_only_summary_note: false,
out: &mut out,
stats: &mut stats,
};
stream_output.print_events(&events).unwrap();
}
let rendered = String::from_utf8(out).unwrap();
let mut lines = rendered.lines();
assert!(lines.next().unwrap().contains("ipdv=n/a"));
assert!(lines.next().unwrap().contains("ipdv=50.0µs"));
}
#[test]
fn continuous_mode_uses_continuous_stats_config() {
let mut collector = StatsCollector::new(stats_config(true));
for seq in 0..5000_u32 {
let sent_at = test_timestamp(Duration::from_micros(seq as u64));
let received_at = test_timestamp(Duration::from_micros(seq as u64 + 1));
collector.process(&ClientEvent::EchoReply {
seq,
remote: test_remote(),
sent_at,
received_at,
rtt: RttSample {
raw: Duration::from_micros(1),
adjusted: None,
effective: SignedDuration::from_nanos(1_000),
},
server_timing: None,
one_way: None,
received_stats: None,
bytes: 64,
packet_meta: PacketMeta::default(),
});
}
assert_eq!(collector.snapshot().rtt.primary.median_ns, None);
}
#[test]
fn finite_stats_memory_warning_reports_only_large_finite_runs() {
for args in [
cli_args(&["--duration", "0", "--interval", "1ms", "127.0.0.1:2112"]),
cli_args(&[
"--duration",
"1000ms",
"--interval",
"1ms",
"127.0.0.1:2112",
]),
] {
assert_eq!(finite_stats_memory_warning(&args), None);
}
for (threshold, size, guidance) in [
(
FINITE_STATS_MEMORY_WARNING_BYTES,
"about 128 MiB",
"bounded-memory long-running tests",
),
(
FINITE_STATS_MEMORY_STRONG_WARNING_BYTES,
"about 512 MiB",
"shortening the run",
),
(
FINITE_STATS_MEMORY_VERY_STRONG_WARNING_BYTES,
"about 1 GiB",
"memory-constrained systems",
),
] {
let expected_probes = threshold.div_ceil(FINITE_STATS_BYTES_PER_PROBE);
let args = cli_args(&[
"--duration",
&format!("{expected_probes}ms"),
"--interval",
"1ms",
"127.0.0.1:2112",
]);
let warning = finite_stats_memory_warning(&args).unwrap();
assert!(warning.contains(size));
assert!(warning.contains(guidance));
}
}
#[test]
fn expected_probe_count_rounds_up() {
assert_eq!(
expected_probe_count(Duration::from_millis(1001), Duration::from_secs(1)),
2
);
}
#[test]
fn estimate_finite_stats_memory_bytes_saturates() {
assert_eq!(estimate_finite_stats_memory_bytes(u64::MAX), u64::MAX);
}
#[cfg(not(feature = "tui"))]
fn run_tui_applet(
_argv: Vec<OsString>,
_shutdown_requested: &AtomicBool,
) -> Result<(), Box<dyn std::error::Error>> {
Err("TUI applet is not available; rebuild with the tui feature".into())
}