use std::{ collections::{BTreeMap, BTreeSet, HashSet, VecDeque}, io::{self, BufRead, Write}, sync::{atomic::AtomicBool, Arc, Mutex}, thread, time::Duration, }; use irtt_client::{ managed::{ BlockingManagedClient, ManagedClientHandle, ManagedCommandApplyError, ManagedEndReason, ManagedEvent, ManagedEventSubscription, ManagedEventTryRecvError, ManagedStatus, ManagedTargetConfig, ManagedTargetEndReason, ManagedTargetOutcome, TargetInstance, }, ClientEvent, }; use super::{ args::ClientArgs, output::{EventRenderStats, OutputConfig}, }; use crate::shared::client::{ expected_probe_count, is_shutdown_requested, parse_stdin_target_set, session::{ drain_managed_events, peer_close_run_error, request_managed_stop_for_peer_close, request_managed_stop_once, should_print_final_summary, ManagedDrainState, }, ManagedRunSetup, STDIN_MAX_DESIRED_TARGETS, STDIN_OUTCOME_HISTORY_LIMIT, }; use irtt_stats::{StatsCollector, StatsConfig}; const MIB: u64 = 1024 * 1024; const GIB: u64 = 1024 * MIB; const FINITE_STATS_MEMORY_WARNING_BYTES: u64 = 128 * MIB; const FINITE_STATS_MEMORY_STRONG_WARNING_BYTES: u64 = 512 * MIB; const FINITE_STATS_MEMORY_VERY_STRONG_WARNING_BYTES: u64 = GIB; const MANAGED_EVENT_WAIT_SLICE: Duration = Duration::from_millis(20); const MAX_STDIN_RECORD_BYTES: usize = 64 * 1024; pub fn run_stream( args: ClientArgs, shutdown_requested: &AtomicBool, ) -> Result<(), Box> { if args.list_columns { print!("{}", OutputConfig::list_columns()); return Ok(()); } let setup = args .prepare() .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?; let multi_target = setup.is_multi_target(); let output_config = OutputConfig::new( args.format, args.columns.as_deref(), args.header, args.verbose, ) .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?; if setup.stdin_controlled { return run_stdin_stream(setup, output_config, shutdown_requested); } let continuous = args.is_continuous(); let target_count = setup.target_count(); if let Some(warning) = finite_stats_memory_warning(&args, target_count) { eprintln!("{warning}"); } if is_shutdown_requested(shutdown_requested) { return Ok(()); } let (owner, mut events) = BlockingManagedClient::start_with_subscription( setup.managed_config(), setup.managed_targets(), )?; let handle = owner.handle(); let mut stdout = io::LineWriter::new(io::stdout().lock()); let mut stream_output = StreamOutput { config: output_config, header_printed: false, print_final_summary: false, show_running_only_summary_note: false, out: &mut stdout, }; let mut stats = setup .targets .iter() .map(|target| { ( target.label.clone(), StatsCollector::new(stats_config(continuous)), ) }) .collect::>(); let mut terminal_targets = HashSet::new(); let mut dropped_events = 0_u64; let mut interrupted = false; let mut stop_requested = false; let mut subscription_closed = false; loop { if is_shutdown_requested(shutdown_requested) { interrupted = true; if request_managed_stop_once(&mut stop_requested) { drop(handle.stop()); } } if request_managed_stop_for_peer_close( continuous, interrupted, handle.status().peer_closed_target_outcomes, &mut stop_requested, ) { drop(handle.stop()); } let drain_state = drain_events( &mut events, &mut stream_output, &mut stats, &mut terminal_targets, &mut dropped_events, )?; match drain_state { ManagedDrainState::Empty => thread::sleep(MANAGED_EVENT_WAIT_SLICE), ManagedDrainState::Closed => subscription_closed = true, ManagedDrainState::BudgetExhausted => {} } if handle.status().final_outcome.is_some() || subscription_closed { break; } } if interrupted { eprintln!("interrupted, closing managed run..."); } let outcome = owner.join()?; drain_final_events( &mut events, &mut stream_output, &mut stats, &mut terminal_targets, &mut dropped_events, )?; if let Some(warning) = dropped_event_warning(dropped_events) { eprintln!("{warning}"); } if outcome.discarded_target_outcomes != 0 { eprintln!( "irtt-rs: warning: {} final target outcomes were discarded", outcome.discarded_target_outcomes ); } for target in outcome.recent_target_outcomes.iter() { if terminal_targets.insert(target.target.clone()) { report_target_failure(target); } } interrupted |= is_shutdown_requested(shutdown_requested); let terminal_error = match &outcome.end_reason { ManagedEndReason::DriverFailed(failure) => { Some(format!("managed driver failed: {failure}")) } _ => peer_close_run_error(continuous, interrupted, outcome.peer_closed_target_outcomes) .or_else(|| { (!interrupted && outcome.successful_target_outcomes == 0 && outcome.failed_target_outcomes > 0) .then(|| { format!( "no managed target completed successfully ({} failed)", outcome.failed_target_outcomes ) }) }), }; stream_output.print_final_summary = should_print_final_summary(continuous, interrupted); stream_output.show_running_only_summary_note = continuous && interrupted && stream_output.print_final_summary; for target in &setup.targets { let target_stats = stats.get(&target.label).expect( "every prepared target must have a stats collector inserted before the run begins", ); if multi_target && stream_output.print_final_summary && stream_output.config.prints_summary() { writeln!(stream_output.out)?; writeln!(stream_output.out, "target: {}", target.label)?; } stream_output.print_summary(target_stats)?; } stream_output.out.flush()?; if let Some(error) = terminal_error { return Err(error.into()); } Ok(()) } #[derive(Clone)] struct StdinTargetSet { revision: u64, targets: Vec, } #[derive(Default)] struct StdinMailboxState { next_revision: u64, latest: Option, eof: bool, fatal: Option, } #[derive(Default)] struct StdinMailbox { state: Mutex, } impl StdinMailbox { fn publish(&self, targets: Vec) { let mut state = self .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); state.next_revision = state.next_revision.saturating_add(1); state.latest = Some(StdinTargetSet { revision: state.next_revision, targets, }); } fn finish(&self) { self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .eof = true; } fn fail(&self, error: String) { let mut state = self .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); state.fatal = Some(error); state.latest = None; } fn latest(&self) -> Option { self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .latest .clone() } fn clear_if_current(&self, revision: u64) { let mut state = self .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if state .latest .as_ref() .is_some_and(|set| set.revision == revision) { state.latest = None; } } fn is_current(&self, revision: u64) -> bool { self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .latest .as_ref() .is_some_and(|set| set.revision == revision) } fn fatal(&self) -> Option { self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .fatal .clone() } fn is_finished(&self) -> bool { self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .eof } } fn read_stdin_record( reader: &mut R, buffer: &mut Vec, ) -> io::Result> { buffer.clear(); loop { let available = reader.fill_buf()?; if available.is_empty() { if buffer.is_empty() { return Ok(None); } if buffer.len() > MAX_STDIN_RECORD_BYTES { return Err(io::Error::new( io::ErrorKind::InvalidData, "stdin target record exceeds the maximum size", )); } return std::str::from_utf8(buffer) .map(str::to_owned) .map(Some) .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "stdin is not UTF-8")); } if buffer.len() == MAX_STDIN_RECORD_BYTES + 1 { if available.first() == Some(&b'\n') && buffer.last() == Some(&b'\r') { reader.consume(1); buffer.pop(); return std::str::from_utf8(buffer) .map(str::to_owned) .map(Some) .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "stdin is not UTF-8")); } return Err(io::Error::new( io::ErrorKind::InvalidData, "stdin target record exceeds the maximum size", )); } if let Some(newline) = available.iter().position(|byte| *byte == b'\n') { let before_newline = &available[..newline]; if buffer.len().saturating_add(before_newline.len()) > MAX_STDIN_RECORD_BYTES + 1 { return Err(io::Error::new( io::ErrorKind::InvalidData, "stdin target record exceeds the maximum size", )); } buffer.extend_from_slice(before_newline); reader.consume(newline + 1); if buffer.last() == Some(&b'\r') { buffer.pop(); } if buffer.len() > MAX_STDIN_RECORD_BYTES { return Err(io::Error::new( io::ErrorKind::InvalidData, "stdin target record exceeds the maximum size", )); } return std::str::from_utf8(buffer) .map(str::to_owned) .map(Some) .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "stdin is not UTF-8")); } let remaining = MAX_STDIN_RECORD_BYTES + 1 - buffer.len(); let take = available.len().min(remaining); buffer.extend_from_slice(&available[..take]); reader.consume(take); if buffer.len() > MAX_STDIN_RECORD_BYTES && buffer.last() != Some(&b'\r') { return Err(io::Error::new( io::ErrorKind::InvalidData, "stdin target record exceeds the maximum size", )); } } } fn read_stdin_target_sets(reader: &mut R, mailbox: &StdinMailbox) { let mut line = 0_u64; let mut record = Vec::with_capacity(MAX_STDIN_RECORD_BYTES + 1); loop { match read_stdin_record(reader, &mut record) { Ok(Some(record)) => { line = line.saturating_add(1); if record.is_empty() { continue; } match parse_stdin_target_set(&record, STDIN_MAX_DESIRED_TARGETS) { Ok(targets) => { mailbox.publish(targets.into_iter().map(|target| target.managed).collect()) } Err(error) => { mailbox.fail(format!("invalid --targets-stdin line {line}: {error}")); return; } } } Ok(None) => { mailbox.finish(); return; } Err(_) => { mailbox.fail(format!("failed to read --targets-stdin line {}", line + 1)); return; } } } } fn spawn_stdin_target_reader(mailbox: Arc) -> io::Result<()> { thread::Builder::new() .name("irtt-targets-stdin".to_owned()) .spawn(move || { let stdin = io::stdin(); read_stdin_target_sets(&mut stdin.lock(), &mailbox); }) .map(|_| ()) } #[derive(Clone)] struct RetryAfterStatus { revision: u64, status: Arc, } fn retry_waits_for_status_change( retry_after: &mut Option, target_set: &StdinTargetSet, status: &Arc, ) -> bool { let Some(retry) = retry_after else { return false; }; if retry.revision != target_set.revision { *retry_after = None; return false; } Arc::ptr_eq(&retry.status, status) } fn retry_after_capacity_rejection( mailbox: &StdinMailbox, target_set: &StdinTargetSet, status: Arc, ) -> Option { mailbox .is_current(target_set.revision) .then_some(RetryAfterStatus { revision: target_set.revision, status, }) } fn apply_latest_stdin_target_set( handle: &ManagedClientHandle, mailbox: &StdinMailbox, retry_after: &mut Option, ) -> Result<(), String> { let Some(target_set) = mailbox.latest() else { *retry_after = None; return Ok(()); }; let status = handle.status(); if retry_waits_for_status_change(retry_after, &target_set, &status) { return Ok(()); } let receipt = handle .update_targets(target_set.targets.clone()) .map_err(|_| "failed to submit --targets-stdin update".to_owned())?; match receipt.blocking_wait() { Ok(_) => { mailbox.clear_if_current(target_set.revision); *retry_after = None; Ok(()) } Err(ManagedCommandApplyError::LiveGenerationLimitExceeded { .. }) => { *retry_after = retry_after_capacity_rejection(mailbox, &target_set, status); Ok(()) } Err(_) => Err("--targets-stdin update was rejected".to_owned()), } } struct BoundedTargetSet { limit: usize, order: VecDeque, members: HashSet, } impl BoundedTargetSet { fn new(limit: usize) -> Self { Self { limit, order: VecDeque::with_capacity(limit), members: HashSet::with_capacity(limit), } } fn insert(&mut self, target: TargetInstance) -> bool { if !self.members.insert(target.clone()) { return false; } if self.order.len() == self.limit { let evicted = self .order .pop_front() .expect("bounded target order is nonempty"); self.members.remove(&evicted); } self.order.push_back(target); true } } fn process_stdin_event( event: ManagedEvent, stream_output: &mut StreamOutput<'_, W>, stats: &mut BTreeMap, terminal_targets: &mut BoundedTargetSet, ) -> io::Result<()> { match event { ManagedEvent::TargetStateChanged { target, .. } => { stats .entry(target) .or_insert_with(|| StatsCollector::new(stats_config(true))); } ManagedEvent::Client { target, event } => { let collector = stats .entry(target.clone()) .or_insert_with(|| StatsCollector::new(stats_config(true))); print_events_with_stats( stream_output, std::slice::from_ref(&event), Some(target.id.as_str()), collector, )?; } ManagedEvent::TargetFinished { outcome } if terminal_targets.insert(outcome.target.clone()) => { report_target_failure(&outcome); } _ => {} } Ok(()) } fn drain_stdin_events( events: &mut ManagedEventSubscription, stream_output: &mut StreamOutput<'_, W>, stats: &mut BTreeMap, terminal_targets: &mut BoundedTargetSet, dropped_events: &mut u64, ) -> io::Result { drain_managed_events(events, dropped_events, |event| { process_stdin_event(event, stream_output, stats, terminal_targets) }) } fn drain_final_stdin_events( events: &mut ManagedEventSubscription, stream_output: &mut StreamOutput<'_, W>, stats: &mut BTreeMap, terminal_targets: &mut BoundedTargetSet, dropped_events: &mut u64, ) -> io::Result<()> { loop { match events.try_recv() { Ok(event) => process_stdin_event(event, stream_output, stats, terminal_targets)?, Err(ManagedEventTryRecvError::Empty | ManagedEventTryRecvError::Closed) => break, Err(ManagedEventTryRecvError::Lagged(count)) => { *dropped_events = dropped_events.saturating_add(count); } } } Ok(()) } fn reconcile_stdin_stats( stats: &mut BTreeMap, status: &ManagedStatus, final_summary_targets: Option<&BTreeSet>, ) { let live = status .targets .iter() .map(|target| target.target.clone()) .collect::>(); stats.retain(|target, _| { live.contains(target) || final_summary_targets.is_some_and(|targets| targets.contains(target)) }); } fn snapshot_stdin_summary_targets( stats: &mut BTreeMap, status: &ManagedStatus, ) -> BTreeSet { let targets = status .targets .iter() .map(|target| target.target.clone()) .collect::>(); for target in &targets { stats .entry(target.clone()) .or_insert_with(|| StatsCollector::new(stats_config(true))); } targets } #[derive(Clone)] enum StdinStop { Interrupted, Eof, Fatal(String), } fn stdin_stop_request( shutdown_requested: &AtomicBool, mailbox: &StdinMailbox, ) -> Option { if is_shutdown_requested(shutdown_requested) { Some(StdinStop::Interrupted) } else if let Some(error) = mailbox.fatal() { Some(StdinStop::Fatal(error)) } else if mailbox.is_finished() { Some(StdinStop::Eof) } else { None } } fn run_stdin_stream( setup: ManagedRunSetup, output_config: OutputConfig, shutdown_requested: &AtomicBool, ) -> Result<(), Box> { if is_shutdown_requested(shutdown_requested) { return Ok(()); } let (owner, mut events) = BlockingManagedClient::start_with_subscription( setup.managed_config(), setup.managed_targets(), )?; let handle = owner.handle(); let mailbox = Arc::new(StdinMailbox::default()); spawn_stdin_target_reader(Arc::clone(&mailbox))?; let mut stdout = io::LineWriter::new(io::stdout().lock()); let mut stream_output = StreamOutput { config: output_config, header_printed: false, print_final_summary: false, show_running_only_summary_note: false, out: &mut stdout, }; let mut stats = BTreeMap::new(); let mut terminal_targets = BoundedTargetSet::new(STDIN_OUTCOME_HISTORY_LIMIT); let mut retry_after = None; let mut dropped_events = 0_u64; let mut stop_requested = false; let mut stop = None; let mut summary_targets = None; let mut subscription_closed = false; loop { let status = handle.status(); if stop.is_none() { if let Some(requested_stop) = stdin_stop_request(shutdown_requested, &mailbox) { summary_targets = Some(snapshot_stdin_summary_targets(&mut stats, &status)); stop = Some(requested_stop); if request_managed_stop_once(&mut stop_requested) { drop(handle.stop()); } } else if let Err(error) = apply_latest_stdin_target_set(&handle, &mailbox, &mut retry_after) { summary_targets = Some(snapshot_stdin_summary_targets(&mut stats, &status)); stop = Some(StdinStop::Fatal(error)); if request_managed_stop_once(&mut stop_requested) { drop(handle.stop()); } } } let drain_state = drain_stdin_events( &mut events, &mut stream_output, &mut stats, &mut terminal_targets, &mut dropped_events, )?; reconcile_stdin_stats(&mut stats, &handle.status(), summary_targets.as_ref()); match drain_state { ManagedDrainState::Empty => thread::sleep(MANAGED_EVENT_WAIT_SLICE), ManagedDrainState::Closed => subscription_closed = true, ManagedDrainState::BudgetExhausted => {} } if handle.status().final_outcome.is_some() || subscription_closed { break; } } let outcome = owner.join()?; drain_final_stdin_events( &mut events, &mut stream_output, &mut stats, &mut terminal_targets, &mut dropped_events, )?; if let Some(warning) = dropped_event_warning(dropped_events) { eprintln!("{warning}"); } for target in outcome.recent_target_outcomes.iter() { if terminal_targets.insert(target.target.clone()) { report_target_failure(target); } } if outcome.discarded_target_outcomes != 0 { eprintln!( "irtt-rs: warning: {} final target outcomes were discarded", outcome.discarded_target_outcomes ); } let print_summary = matches!(stop, Some(StdinStop::Interrupted | StdinStop::Eof)); stream_output.print_final_summary = print_summary; stream_output.show_running_only_summary_note = print_summary; if let Some(summary_targets) = summary_targets { let multi_target = summary_targets.len() > 1; for target in summary_targets { let stats = stats .get(&target) .expect("summary snapshot retains every selected target collector"); if multi_target && stream_output.config.prints_summary() { writeln!(stream_output.out)?; writeln!( stream_output.out, "target: {} (generation {})", target.id, target.generation )?; } stream_output.print_summary(stats)?; } } stream_output.out.flush()?; match stop { Some(StdinStop::Fatal(error)) => Err(error.into()), _ => match outcome.end_reason { ManagedEndReason::DriverFailed(failure) => { Err(format!("managed driver failed: {failure}").into()) } _ => Ok(()), }, } } fn process_event( event: ManagedEvent, stream_output: &mut StreamOutput<'_, W>, stats: &mut BTreeMap, terminal_targets: &mut HashSet, ) -> io::Result<()> { match event { ManagedEvent::Client { target, event } => { let collector = stats .entry(target.id.as_str().to_owned()) .or_insert_with(|| StatsCollector::new(stats_config(false))); print_events_with_stats( stream_output, std::slice::from_ref(&event), Some(target.id.as_str()), collector, )?; } ManagedEvent::TargetFinished { outcome } => { terminal_targets.insert(outcome.target.clone()); report_target_failure(&outcome); } _ => {} } Ok(()) } fn drain_events( events: &mut ManagedEventSubscription, stream_output: &mut StreamOutput<'_, W>, stats: &mut BTreeMap, terminal_targets: &mut HashSet, dropped_events: &mut u64, ) -> io::Result { drain_managed_events(events, dropped_events, |event| { process_event(event, stream_output, stats, terminal_targets) }) } fn drain_final_events( events: &mut ManagedEventSubscription, stream_output: &mut StreamOutput<'_, W>, stats: &mut BTreeMap, terminal_targets: &mut HashSet, dropped_events: &mut u64, ) -> io::Result<()> { loop { match events.try_recv() { Ok(event) => process_event(event, stream_output, stats, terminal_targets)?, Err(ManagedEventTryRecvError::Empty | ManagedEventTryRecvError::Closed) => break, Err(ManagedEventTryRecvError::Lagged(count)) => { *dropped_events = dropped_events.saturating_add(count); } } } Ok(()) } fn report_target_failure(target: &ManagedTargetOutcome) { for message in target_failure_messages(target) { eprintln!("{message}"); } } fn target_failure_messages(target: &ManagedTargetOutcome) -> Vec { let mut messages = Vec::with_capacity(2); if let ManagedTargetEndReason::Failed(failure) = &target.end_reason { messages.push(format!( "irtt-rs: target {} failed ({} {}): {}", target.target.id, failure.phase, failure.kind, failure.message )); } if let Some(failure) = &target.cleanup_failure { messages.push(format!( "irtt-rs: target {} cleanup failed ({} {}): {}", target.target.id, failure.phase, failure.kind, failure.message )); } messages } fn dropped_event_warning(dropped_events: u64) -> Option { (dropped_events > 0).then(|| format!("irtt-rs: warning: dropped {dropped_events} managed run event{}; output and statistics may be incomplete", if dropped_events == 1 { "" } else { "s" })) } struct StreamOutput<'a, W: Write> { config: OutputConfig, header_printed: bool, print_final_summary: bool, show_running_only_summary_note: bool, out: &'a mut W, } impl StreamOutput<'_, W> { fn print_events( &mut self, events: &[ClientEvent], target: Option<&str>, stats_updates: &[EventRenderStats], ) -> io::Result<()> { self.print_header()?; for (event, stats_update) in events.iter().zip(stats_updates) { if let Some(line) = self.config.render_event(event, target, Some(stats_update)) { writeln!(self.out, "{line}")?; } } Ok(()) } fn print_header(&mut self) -> io::Result<()> { if self.header_printed { return Ok(()); } self.header_printed = true; if let Some(header) = self.config.render_header() { writeln!(self.out, "{header}")?; } Ok(()) } fn print_summary(&mut self, stats: &StatsCollector) -> io::Result<()> { if self.print_final_summary && self.config.prints_summary() { write!( self.out, "{}", crate::cmd::client::summary::format_summary_with_options( &stats.snapshot(), crate::cmd::client::summary::SummaryFormatOptions { verbose: self.config.summary_verbose(), show_running_only_note: self.show_running_only_summary_note } ) )?; } Ok(()) } } fn print_events_with_stats( stream_output: &mut StreamOutput<'_, W>, events: &[ClientEvent], target: Option<&str>, stats: &mut StatsCollector, ) -> io::Result<()> { let updates = events .iter() .map(|event| EventRenderStats::from(stats.process(event))) .collect::>(); stream_output.print_events(events, target, &updates) } fn stats_config(continuous: bool) -> StatsConfig { if continuous { StatsConfig::continuous() } else { StatsConfig::finite() } } fn finite_stats_memory_warning(args: &ClientArgs, target_count: usize) -> Option { if args.is_continuous() || args.duration.is_zero() { return None; } let target_count = u64::try_from(target_count).unwrap_or(u64::MAX); let total_probe_count = expected_probe_count(args.duration, args.interval).saturating_mul(target_count); // Ask the stats crate what the configuration this run will actually use // retains; the CLI owns the probe count and the thresholds, not the // retention model. let estimated_bytes = stats_config(false).estimated_retained_bytes(total_probe_count); if estimated_bytes < FINITE_STATS_MEMORY_WARNING_BYTES { return None; } let formatted = if estimated_bytes >= GIB { format!("{} GiB", estimated_bytes.saturating_add(GIB / 2) / GIB) } else { format!("{} MiB", estimated_bytes.saturating_add(MIB / 2) / MIB) }; 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(test)] mod tests { use super::*; use irtt_client::managed::{ ManagedLifecycle, ManagedTargetFailure, ManagedTargetFailureKind, ManagedTargetFailurePhase, ManagedTargetLifecycle, ManagedTargetStatus, TargetId, }; use std::{io::BufReader, io::Cursor, sync::Arc}; #[test] fn stdin_record_reader_removes_only_line_framing() { let mut reader = BufReader::with_capacity(2, Cursor::new(b" first \r\n\nthird\r")); let mut buffer = Vec::with_capacity(MAX_STDIN_RECORD_BYTES + 1); assert_eq!( read_stdin_record(&mut reader, &mut buffer).unwrap(), Some(" first ".to_owned()) ); assert_eq!( read_stdin_record(&mut reader, &mut buffer).unwrap(), Some(String::new()) ); assert_eq!( read_stdin_record(&mut reader, &mut buffer).unwrap(), Some("third\r".to_owned()) ); assert_eq!(read_stdin_record(&mut reader, &mut buffer).unwrap(), None); } #[test] fn stdin_record_reader_rejects_over_limit_unterminated_record() { let input = vec![b'x'; MAX_STDIN_RECORD_BYTES + 1]; let mut reader = BufReader::with_capacity(17, Cursor::new(input)); let mut buffer = Vec::with_capacity(MAX_STDIN_RECORD_BYTES + 1); let error = read_stdin_record(&mut reader, &mut buffer).unwrap_err(); assert_eq!(error.kind(), io::ErrorKind::InvalidData); assert!(error.to_string().contains("maximum size")); assert!(buffer.len() <= MAX_STDIN_RECORD_BYTES + 1); } #[test] fn newer_mailbox_value_survives_an_earlier_submitted_acknowledgement() { let mailbox = StdinMailbox::default(); mailbox.publish(vec![ManagedTargetConfig::new("a", "127.0.0.1:2112")]); let submitted = mailbox.latest().unwrap(); mailbox.publish(vec![ManagedTargetConfig::new("b", "127.0.0.1:2113")]); let pending = mailbox.latest().unwrap(); mailbox.clear_if_current(submitted.revision); assert_eq!(mailbox.latest().unwrap().revision, pending.revision); } #[test] fn rejected_update_retries_after_status_publishes_during_acknowledgement() { let mailbox = StdinMailbox::default(); mailbox.publish(Vec::new()); let target_set = mailbox.latest().unwrap(); let saturated = Arc::new(status_with_targets(&[TargetInstance { id: TargetId::from("old"), generation: 1, }])); let capacity_freed = Arc::new(status_with_targets(&[])); // Capacity becomes available after the transaction rejected but before // its acknowledgement reaches the controller. The rejection path must // retain the pre-submission status, not this later snapshot. let mut retry_after = retry_after_capacity_rejection(&mailbox, &target_set, Arc::clone(&saturated)); assert!(retry_waits_for_status_change( &mut retry_after, &target_set, &saturated )); assert!(!retry_waits_for_status_change( &mut retry_after, &target_set, &capacity_freed )); } #[test] fn later_desired_set_supersedes_a_capacity_rejected_candidate_before_retry() { let mailbox = StdinMailbox::default(); mailbox.publish(vec![ManagedTargetConfig::new("a", "127.0.0.1:2112")]); let applied = mailbox.latest().unwrap(); mailbox.clear_if_current(applied.revision); mailbox.publish(vec![ManagedTargetConfig::new("b", "127.0.0.1:2113")]); let rejected = mailbox.latest().unwrap(); let saturated = Arc::new(status_with_targets(&[TargetInstance { id: TargetId::from("old"), generation: 1, }])); let mut retry_after = retry_after_capacity_rejection(&mailbox, &rejected, saturated); mailbox.publish(vec![ManagedTargetConfig::new("c", "127.0.0.1:2114")]); let replacement = mailbox.latest().unwrap(); assert!(!retry_waits_for_status_change( &mut retry_after, &replacement, &Arc::new(status_with_targets(&[TargetInstance { id: TargetId::from("old"), generation: 1, }])) )); assert!(retry_after.is_none()); assert_eq!(replacement.targets[0].id.as_str(), "c"); mailbox.clear_if_current(replacement.revision); assert!(mailbox.latest().is_none()); } fn status_with_targets(targets: &[TargetInstance]) -> ManagedStatus { ManagedStatus { lifecycle: ManagedLifecycle::Running, stop_requested: false, applied_command_sequence: 0, desired_target_count: targets.len(), connecting_target_count: 0, opening_target_count: 0, active_target_count: 0, draining_target_count: 0, closing_target_count: 0, terminal_target_count: 0, total_target_outcomes: 0, successful_target_outcomes: 0, failed_target_outcomes: 0, peer_closed_target_outcomes: 0, discarded_target_outcomes: 0, targets: Arc::from( targets .iter() .cloned() .map(|target| ManagedTargetStatus { target, desired: true, lifecycle: ManagedTargetLifecycle::Active, server_addr: Arc::from("127.0.0.1:2112"), remote: None, }) .collect::>() .into_boxed_slice(), ), recent_target_outcomes: Arc::from([]), final_outcome: None, } } #[test] fn dynamic_accounting_prunes_retired_generations_and_retains_only_shutdown_snapshot() { let retired = TargetInstance { id: TargetId::from("edge"), generation: 1, }; let replacement = TargetInstance { id: TargetId::from("edge"), generation: 2, }; let live_status = status_with_targets(std::slice::from_ref(&replacement)); let empty_status = status_with_targets(&[]); let mut stats = BTreeMap::from([ (retired.clone(), StatsCollector::new(stats_config(true))), (replacement.clone(), StatsCollector::new(stats_config(true))), ]); reconcile_stdin_stats(&mut stats, &live_status, None); assert_eq!(stats.keys().cloned().collect::>(), [replacement]); let summary_targets = snapshot_stdin_summary_targets(&mut stats, &live_status); stats.insert(retired, StatsCollector::new(stats_config(true))); reconcile_stdin_stats(&mut stats, &empty_status, Some(&summary_targets)); assert_eq!( stats.keys().cloned().collect::>(), [TargetInstance { id: TargetId::from("edge"), generation: 2, }] ); } #[test] fn terminal_deduplication_is_bounded_under_target_churn() { let mut terminal_targets = BoundedTargetSet::new(2); let target = |generation| TargetInstance { id: TargetId::from("edge"), generation, }; assert!(terminal_targets.insert(target(1))); assert!(terminal_targets.insert(target(2))); assert!(terminal_targets.insert(target(3))); assert!(terminal_targets.insert(target(1))); assert!(!terminal_targets.insert(target(3))); } fn long_finite_args() -> ClientArgs { use clap::Parser; ClientArgs::try_parse_from([ "irtt-client", "--duration", "4200s", "--interval", "100ms", "127.0.0.1:2112", ]) .unwrap() } /// Smallest target count whose aggregate estimate reaches the warning /// threshold, derived from the stats crate's own estimate rather than from /// any bytes-per-probe assumption held here. fn first_warning_target_count(args: &ClientArgs) -> usize { let probes = expected_probe_count(args.duration, args.interval); (1..=1024) .find(|targets| { let total = probes.saturating_mul(*targets as u64); stats_config(false).estimated_retained_bytes(total) >= FINITE_STATS_MEMORY_WARNING_BYTES }) .expect("a long finite run should cross the warning threshold within 1024 targets") } #[test] fn finite_stats_memory_warning_accounts_for_target_count() { let args = long_finite_args(); let crossing = first_warning_target_count(&args); assert!( crossing > 1, "this run should need more than one target to cross the threshold, \ otherwise the test proves nothing about target-count accounting" ); assert!( finite_stats_memory_warning(&args, crossing - 1).is_none(), "an aggregate estimate below the threshold should not warn" ); assert!( finite_stats_memory_warning(&args, crossing).is_some(), "an aggregate estimate at the threshold should warn" ); } #[test] fn finite_stats_memory_warning_saturates_target_count_multiplication() { assert!(finite_stats_memory_warning(&long_finite_args(), usize::MAX).is_some()); } #[test] fn target_failure_diagnostics_include_primary_and_cleanup_failures() { let target = ManagedTargetOutcome { target: TargetInstance { id: TargetId::from("edge"), generation: 1, }, server_addr: Arc::from("127.0.0.1:2112"), remote: None, end_reason: ManagedTargetEndReason::Failed(ManagedTargetFailure { phase: ManagedTargetFailurePhase::Receiving, kind: ManagedTargetFailureKind::Protocol, message: "primary failure".into(), }), packets_sent: 0, replies_received: 0, duplicates: 0, late: 0, warning_events: 0, cleanup_failure: Some(ManagedTargetFailure { phase: ManagedTargetFailurePhase::Closing, kind: ManagedTargetFailureKind::Socket, message: "cleanup failure".into(), }), }; assert_eq!( target_failure_messages(&target), [ "irtt-rs: target edge failed (receiving protocol): primary failure", "irtt-rs: target edge cleanup failed (closing socket): cleanup failure", ] ); } }