diff --git a/crates/irtt-app/Cargo.toml b/crates/irtt-app/Cargo.toml index 64bd9c7..10552d2 100644 --- a/crates/irtt-app/Cargo.toml +++ b/crates/irtt-app/Cargo.toml @@ -42,7 +42,7 @@ doc = false [features] default = ["client", "server", "tui"] -client = ["dep:irtt-client", "irtt-client/ancillary", "irtt-client/tokio", "dep:irtt-proto", "dep:irtt-stats"] +client = ["dep:tokio", "dep:irtt-client", "irtt-client/ancillary", "irtt-client/tokio", "dep:irtt-proto", "dep:irtt-stats"] server = ["dep:irtt-server", "dep:tokio"] tui =["dep:irtt-client", "irtt-client/ancillary", "irtt-client/tokio", "dep:irtt-proto", "dep:irtt-stats", "dep:ratatui", "dep:crossterm"] @@ -55,7 +55,7 @@ irtt-proto = { workspace = true, optional = true } irtt-server = { workspace = true, optional = true } irtt-stats = { workspace = true, optional = true } ratatui = { version = "0.30.2", optional = true } -tokio = { version = "1.53.2", default-features = false, features = ["rt", "time"], optional = true } +tokio = { version = "1.53.2", default-features = false, features = ["rt", "time", "sync", "macros"], optional = true } [dev-dependencies] irtt-server.workspace = true diff --git a/crates/irtt-app/src/cmd/client/mod.rs b/crates/irtt-app/src/cmd/client/mod.rs index 44bbb87..8de33e2 100644 --- a/crates/irtt-app/src/cmd/client/mod.rs +++ b/crates/irtt-app/src/cmd/client/mod.rs @@ -2,6 +2,7 @@ pub mod args; pub mod output; pub mod run; pub mod summary; +mod worker; pub use args::{ClientArgs, HeaderMode, OutputFormat}; pub use run::run_stream; diff --git a/crates/irtt-app/src/cmd/client/run.rs b/crates/irtt-app/src/cmd/client/run.rs index f6ac03c..c9583c1 100644 --- a/crates/irtt-app/src/cmd/client/run.rs +++ b/crates/irtt-app/src/cmd/client/run.rs @@ -1,16 +1,15 @@ use std::{ collections::{BTreeMap, BTreeSet, HashSet, VecDeque}, io::{self, BufRead, Write}, - sync::{atomic::AtomicBool, Arc, Mutex}, + sync::Arc, thread, - time::Duration, }; use irtt_client::{ managed::{ - BlockingManagedClient, ManagedClientHandle, ManagedCommandApplyError, ManagedEndReason, - ManagedEvent, ManagedEventSubscription, ManagedEventTryRecvError, ManagedStatus, - ManagedTargetConfig, ManagedTargetEndReason, ManagedTargetOutcome, TargetInstance, + ManagedClient, ManagedCommandApplyError, ManagedEndReason, ManagedEvent, + ManagedEventSubscription, ManagedEventTryRecvError, ManagedStatus, ManagedTargetConfig, + ManagedTargetEndReason, ManagedTargetOutcome, TargetInstance, }, ClientEvent, }; @@ -18,13 +17,14 @@ use irtt_client::{ use super::{ args::ClientArgs, output::{EventRenderStats, OutputConfig}, + worker::ManagedWorker, }; use crate::shared::client::{ - expected_probe_count, is_shutdown_requested, parse_stdin_target_set, + expected_probe_count, 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, + peer_close_run_error, request_managed_stop_for_peer_close, request_managed_stop_once, + should_print_final_summary, }, ManagedRunSetup, STDIN_MAX_DESIRED_TARGETS, STDIN_OUTCOME_HISTORY_LIMIT, }; @@ -40,12 +40,11 @@ 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( +pub async fn run_stream( args: ClientArgs, - shutdown_requested: &AtomicBool, + mut shutdown: tokio::sync::watch::Receiver, ) -> Result<(), Box> { if args.list_columns { print!("{}", OutputConfig::list_columns()); @@ -63,7 +62,7 @@ pub fn run_stream( ) .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?; if setup.stdin_controlled { - return run_stdin_stream(setup, output_config, shutdown_requested); + return run_stdin_stream(setup, output_config, shutdown).await; } let continuous = args.is_continuous(); let target_count = setup.target_count(); @@ -71,14 +70,13 @@ pub fn run_stream( if let Some(warning) = finite_stats_memory_warning(&args, target_count) { eprintln!("{warning}"); } - if is_shutdown_requested(shutdown_requested) { + if *shutdown.borrow() { return Ok(()); } - let (owner, mut events) = BlockingManagedClient::start_with_subscription( - setup.managed_config(), - setup.managed_targets(), - )?; - let handle = owner.handle(); + let (task, handle) = ManagedClient::task(setup.managed_config(), setup.managed_targets())?; + let mut events = handle.subscribe()?; + let mut status = handle.subscribe_status(); + let owner = ManagedWorker::start(task, handle.clone())?; let mut stdout = io::LineWriter::new(io::stdout().lock()); let mut stream_output = StreamOutput { config: output_config, @@ -102,9 +100,9 @@ pub fn run_stream( 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) { + let snapshot = status.borrow_and_update().clone(); + if *shutdown.borrow() { interrupted = true; if request_managed_stop_once(&mut stop_requested) { drop(handle.stop()); @@ -113,38 +111,33 @@ pub fn run_stream( if request_managed_stop_for_peer_close( continuous, interrupted, - handle.status().peer_closed_target_outcomes, + snapshot.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 { + if snapshot.final_outcome.is_some() { break; } + tokio::select! { + _ = shutdown.changed(), if !interrupted => {} + changed = status.changed() => { if changed.is_err() { break; } } + event = events.recv() => match event { + Ok(event) => process_event(event, &mut stream_output, &mut stats, &mut terminal_targets)?, + Err(tokio::sync::broadcast::error::RecvError::Lagged(count)) => { + dropped_events = dropped_events.saturating_add(count); + } + Err(tokio::sync::broadcast::error::RecvError::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, - )?; + let outcome = owner.join().await?; + drain_final_events(&mut events, &mut dropped_events, |event| { + process_event(event, &mut stream_output, &mut stats, &mut terminal_targets) + })?; if let Some(warning) = dropped_event_warning(dropped_events) { eprintln!("{warning}"); } @@ -159,7 +152,7 @@ pub fn run_stream( report_target_failure(target); } } - interrupted |= is_shutdown_requested(shutdown_requested); + interrupted |= *shutdown.borrow(); let terminal_error = match &outcome.end_reason { ManagedEndReason::DriverFailed(failure) => { Some(format!("managed driver failed: {failure}")) @@ -201,98 +194,9 @@ pub fn run_stream( } #[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 - } +enum StdinUpdate { + Targets(Vec), + Stop(StdinStop), } fn read_stdin_record( @@ -371,7 +275,10 @@ fn read_stdin_record( } } -fn read_stdin_target_sets(reader: &mut R, mailbox: &StdinMailbox) { +fn read_stdin_target_sets( + reader: &mut R, + updates: &tokio::sync::watch::Sender, +) { let mut line = 0_u64; let mut record = Vec::with_capacity(MAX_STDIN_RECORD_BYTES + 1); loop { @@ -383,101 +290,43 @@ fn read_stdin_target_sets(reader: &mut R, mailbox: &StdinMailbox) { } match parse_stdin_target_set(&record, STDIN_MAX_DESIRED_TARGETS) { Ok(targets) => { - mailbox.publish(targets.into_iter().map(|target| target.managed).collect()) + updates.send_replace(StdinUpdate::Targets( + targets.into_iter().map(|target| target.managed).collect(), + )); } Err(error) => { - mailbox.fail(format!("invalid --targets-stdin line {line}: {error}")); + updates.send_replace(StdinUpdate::Stop(StdinStop::Fatal(format!( + "invalid --targets-stdin line {line}: {error}" + )))); return; } } } Ok(None) => { - mailbox.finish(); + updates.send_replace(StdinUpdate::Stop(StdinStop::Eof)); return; } Err(_) => { - mailbox.fail(format!("failed to read --targets-stdin line {}", line + 1)); + updates.send_replace(StdinUpdate::Stop(StdinStop::Fatal(format!( + "failed to read --targets-stdin line {}", + line + 1 + )))); return; } } } } -fn spawn_stdin_target_reader(mailbox: Arc) -> io::Result<()> { +fn spawn_stdin_target_reader(updates: tokio::sync::watch::Sender) -> 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); + read_stdin_target_sets(&mut stdin.lock(), &updates); }) .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, @@ -542,37 +391,6 @@ fn process_stdin_event( 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, @@ -613,36 +431,20 @@ enum StdinStop { 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( +async fn run_stdin_stream( setup: ManagedRunSetup, output_config: OutputConfig, - shutdown_requested: &AtomicBool, + mut shutdown: tokio::sync::watch::Receiver, ) -> Result<(), Box> { - if is_shutdown_requested(shutdown_requested) { + if *shutdown.borrow() { 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 (task, handle) = ManagedClient::task(setup.managed_config(), setup.managed_targets())?; + let mut events = handle.subscribe()?; + let mut status = handle.subscribe_status(); + let owner = ManagedWorker::start(task, handle.clone())?; + let (updates, mut stdin) = tokio::sync::watch::channel(StdinUpdate::Targets(Vec::new())); + spawn_stdin_target_reader(updates)?; let mut stdout = io::LineWriter::new(io::stdout().lock()); let mut stream_output = StreamOutput { @@ -654,59 +456,103 @@ fn run_stdin_stream( }; let mut stats = BTreeMap::new(); let mut terminal_targets = BoundedTargetSet::new(STDIN_OUTCOME_HISTORY_LIMIT); + let mut stdin_changed = false; + let mut desired = None; + let mut pending = None; + let mut submitted = Vec::new(); 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(); + let snapshot = status.borrow_and_update().clone(); 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()); + if *shutdown.borrow() { + stop = Some(StdinStop::Interrupted); + } else if stdin_changed || stdin.has_changed().unwrap_or(true) { + stdin_changed = false; + match stdin.borrow_and_update().clone() { + StdinUpdate::Targets(targets) => { + desired = Some(targets); + retry_after = None; + } + StdinUpdate::Stop(reason) => stop = Some(reason), } - } else if let Err(error) = - apply_latest_stdin_target_set(&handle, &mailbox, &mut retry_after) + } + if stop.is_none() + && pending.is_none() + && !retry_after + .as_ref() + .is_some_and(|previous| Arc::ptr_eq(previous, &snapshot)) { - summary_targets = Some(snapshot_stdin_summary_targets(&mut stats, &status)); - stop = Some(StdinStop::Fatal(error)); + if let Some(targets) = desired.take() { + submitted = targets; + match handle.update_targets(submitted.clone()) { + Ok(receipt) => { + pending = Some(receipt); + retry_after = Some(snapshot.clone()); + } + Err(_) => { + stop = Some(StdinStop::Fatal( + "failed to submit --targets-stdin update".to_owned(), + )) + } + } + } + } + if stop.is_some() { + summary_targets = Some(snapshot_stdin_summary_targets(&mut stats, &snapshot)); 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 { + reconcile_stdin_stats(&mut stats, &snapshot, summary_targets.as_ref()); + if snapshot.final_outcome.is_some() { break; } + tokio::select! { + _ = shutdown.changed(), if stop.is_none() => {} + _ = stdin.changed(), if stop.is_none() => { stdin_changed = true; } + changed = status.changed() => { if changed.is_err() { break; } } + result = async { pending.as_mut().expect("guarded pending receipt").await }, if pending.is_some() && stop.is_none() => { + pending = None; + match result { + Ok(_) => retry_after = None, + Err(ManagedCommandApplyError::LiveGenerationLimitExceeded { .. }) => { + // A newer desired set supersedes this rejected set. Otherwise, + // retry only after the driver's status advances from submission. + if desired.is_none() { + desired = Some(std::mem::take(&mut submitted)); + } else { + retry_after = None; + } + } + Err(_) => { + stop = Some(StdinStop::Fatal("--targets-stdin update was rejected".to_owned())); + summary_targets = Some(snapshot_stdin_summary_targets(&mut stats, &status.borrow())); + if request_managed_stop_once(&mut stop_requested) { drop(handle.stop()); } + } + } + } + event = events.recv() => match event { + Ok(event) => { + process_stdin_event(event, &mut stream_output, &mut stats, &mut terminal_targets)?; + reconcile_stdin_stats(&mut stats, &status.borrow(), summary_targets.as_ref()); + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(count)) => { + dropped_events = dropped_events.saturating_add(count); + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => break, + } + } } - let outcome = owner.join()?; - drain_final_stdin_events( - &mut events, - &mut stream_output, - &mut stats, - &mut terminal_targets, - &mut dropped_events, - )?; + let outcome = owner.join().await?; + drain_final_events(&mut events, &mut dropped_events, |event| { + process_stdin_event(event, &mut stream_output, &mut stats, &mut terminal_targets) + })?; if let Some(warning) = dropped_event_warning(dropped_events) { eprintln!("{warning}"); } @@ -782,28 +628,14 @@ fn process_event( Ok(()) } -fn drain_events( +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 { - 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<()> { + mut process: impl FnMut(ManagedEvent) -> Result<(), E>, +) -> Result<(), E> { loop { match events.try_recv() { - Ok(event) => process_event(event, stream_output, stats, terminal_targets)?, + Ok(event) => process(event)?, Err(ManagedEventTryRecvError::Empty | ManagedEventTryRecvError::Closed) => break, Err(ManagedEventTryRecvError::Lagged(count)) => { *dropped_events = dropped_events.saturating_add(count); diff --git a/crates/irtt-app/src/cmd/client/worker.rs b/crates/irtt-app/src/cmd/client/worker.rs new file mode 100644 index 0000000..d9d646c --- /dev/null +++ b/crates/irtt-app/src/cmd/client/worker.rs @@ -0,0 +1,133 @@ +use std::{io, thread, time::Duration}; + +use irtt_client::managed::{ManagedClientHandle, ManagedClientTask, ManagedOutcome}; +use tokio::sync::oneshot; + +/// Keeps measurement scheduling separate from frontend output and statistics. +/// Drop also stops and joins on frontend errors or cancellation. +pub(super) struct ManagedWorker { + handle: ManagedClientHandle, + thread: Option>, + completion: oneshot::Receiver>, +} + +impl ManagedWorker { + pub(super) fn start(task: ManagedClientTask, handle: ManagedClientHandle) -> io::Result { + let (completed, completion) = oneshot::channel(); + let thread = thread::Builder::new() + .name("irtt-managed".to_owned()) + .spawn(move || { + // Construct and destroy the runtime on its owning thread, including + // startup failures; never drop a runtime inside the frontend executor. + let outcome = (|| { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + runtime.block_on(task) + })); + runtime.shutdown_timeout(Duration::from_millis(100)); + match outcome { + Ok(outcome) => Ok(outcome), + Err(panic) => std::panic::resume_unwind(panic), + } + })(); + let _ = completed.send(outcome); + })?; + Ok(Self { + handle, + thread: Some(thread), + completion, + }) + } + + pub(super) async fn join(mut self) -> Result> { + let outcome = (&mut self.completion).await; + self.join_thread()?; + Ok(outcome + .map_err(|_| io::Error::other("managed worker completed without an outcome"))??) + } + + fn join_thread(&mut self) -> io::Result<()> { + self.thread + .take() + .expect("managed worker is joined once") + .join() + .map_err(|_| io::Error::other("managed worker thread panicked")) + } +} + +impl Drop for ManagedWorker { + fn drop(&mut self) { + if self.thread.is_some() { + drop(self.handle.stop()); + let _ = self.join_thread(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use irtt_client::managed::{ + ManagedClient, ManagedClientConfig, ManagedCompletionPolicy, ManagedEndReason, + }; + + // The frontend executor must be able to stop a worker while awaiting its + // completion; a blocking join here would deadlock this current-thread runtime. + #[test] + fn completion_leaves_frontend_free_to_stop_and_preserves_outcome() { + let (task, handle) = ManagedClient::task( + ManagedClientConfig { + completion: ManagedCompletionPolicy::ExplicitStop, + ..Default::default() + }, + vec![], + ) + .unwrap(); + let worker = ManagedWorker::start(task, handle.clone()).unwrap(); + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + let (outcome, ()) = tokio::time::timeout(Duration::from_secs(2), async { + tokio::join!(worker.join(), handle.stop()) + }) + .await + .expect("worker completion must allow frontend stop"); + let outcome = outcome.unwrap(); + assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested); + assert_eq!(handle.status().final_outcome.as_deref(), Some(&outcome)); + }); + } + + // Cancelling the frontend's completion wait must stop and join the worker, + // rather than detaching an ongoing measurement task. + #[test] + fn cancelling_completion_stops_and_joins_worker() { + use std::{ + future::Future, + task::{Context, Waker}, + }; + let (task, handle) = ManagedClient::task( + ManagedClientConfig { + completion: ManagedCompletionPolicy::ExplicitStop, + ..Default::default() + }, + vec![], + ) + .unwrap(); + let worker = ManagedWorker::start(task, handle.clone()).unwrap(); + let mut completion = Box::pin(worker.join()); + assert!(completion + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + .is_pending()); + drop(completion); + assert_eq!( + handle.status().final_outcome.as_ref().unwrap().end_reason, + ManagedEndReason::StopRequested + ); + } +} diff --git a/crates/irtt-app/src/entry.rs b/crates/irtt-app/src/entry.rs index 10fac74..6280c95 100644 --- a/crates/irtt-app/src/entry.rs +++ b/crates/irtt-app/src/entry.rs @@ -1,9 +1,9 @@ -#[cfg(any(feature = "client", feature = "server", feature = "tui"))] +#[cfg(any(feature = "server", feature = "tui"))] use std::sync::{atomic::AtomicBool, Arc}; use std::{env, ffi::OsString, process::ExitCode}; use crate::applet::{dispatch_from_argv, AppletDispatch, RequestedApplet}; -#[cfg(any(feature = "client", feature = "server", feature = "tui"))] +#[cfg(any(feature = "server", feature = "tui"))] use crate::signal::install_signal_handler; /// Entry point for the `irtt-rs` multicall dispatcher binary. @@ -75,9 +75,9 @@ fn run_dispatcher_from_env() -> Result<(), Box> { } } -/// Installs the shutdown signal handler shared by every applet and hands the +/// Installs the shutdown signal handler for synchronous applets and hands the /// resulting flag to `f`. -#[cfg(any(feature = "client", feature = "server", feature = "tui"))] +#[cfg(any(feature = "server", feature = "tui"))] fn with_shutdown_flag( f: impl FnOnce(&AtomicBool) -> Result<(), Box>, ) -> Result<(), Box> { @@ -91,10 +91,13 @@ fn with_shutdown_flag( fn run_client_applet(argv: Vec) -> Result<(), Box> { use clap::Parser; - with_shutdown_flag(|shutdown_requested| { - let args = crate::cmd::client::ClientArgs::parse_from(argv); - crate::cmd::client::run_stream(args, shutdown_requested) - }) + let shutdown = crate::signal::install_async_signal_handler() + .map_err(|err| format!("failed to install signal handler: {err}"))?; + let args = crate::cmd::client::ClientArgs::parse_from(argv); + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()? + .block_on(crate::cmd::client::run_stream(args, shutdown)) } #[cfg(not(feature = "client"))] diff --git a/crates/irtt-app/src/shared/client/session.rs b/crates/irtt-app/src/shared/client/session.rs index 26cde72..5871771 100644 --- a/crates/irtt-app/src/shared/client/session.rs +++ b/crates/irtt-app/src/shared/client/session.rs @@ -1,12 +1,12 @@ use std::sync::atomic::{AtomicBool, Ordering}; -#[cfg(any(feature = "client", feature = "tui"))] +#[cfg(feature = "tui")] use irtt_client::managed::{ManagedEvent, ManagedEventSubscription, ManagedEventTryRecvError}; -#[cfg(any(feature = "client", feature = "tui"))] +#[cfg(feature = "tui")] pub(crate) const MANAGED_EVENT_WORK_BUDGET: usize = 128; -#[cfg(any(feature = "client", feature = "tui"))] +#[cfg(feature = "tui")] #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(crate) enum ManagedDrainState { Empty, @@ -16,7 +16,7 @@ pub(crate) enum ManagedDrainState { /// Drain a bounded amount of lossy presentation work before returning to the /// frontend control loop. -#[cfg(any(feature = "client", feature = "tui"))] +#[cfg(feature = "tui")] pub(crate) fn drain_managed_events( events: &mut ManagedEventSubscription, dropped_events: &mut u64, diff --git a/crates/irtt-app/src/signal.rs b/crates/irtt-app/src/signal.rs index 9bc6249..8c4e0a2 100644 --- a/crates/irtt-app/src/signal.rs +++ b/crates/irtt-app/src/signal.rs @@ -1,10 +1,22 @@ +#[cfg(any(feature = "server", feature = "tui"))] use std::sync::{ atomic::{AtomicBool, Ordering}, Arc, }; +#[cfg(any(feature = "server", feature = "tui"))] pub fn install_signal_handler(shutdown_requested: Arc) -> Result<(), ctrlc::Error> { ctrlc::set_handler(move || { shutdown_requested.store(true, Ordering::Relaxed); }) } + +/// Preserve ctrlc's SIGINT/SIGTERM policy while waking the async client. +#[cfg(feature = "client")] +pub fn install_async_signal_handler() -> Result, ctrlc::Error> { + let (shutdown, receiver) = tokio::sync::watch::channel(false); + ctrlc::set_handler(move || { + shutdown.send_replace(true); + })?; + Ok(receiver) +} diff --git a/crates/irtt-app/tests/cli.rs b/crates/irtt-app/tests/cli.rs index a66f492..fb453f7 100644 --- a/crates/irtt-app/tests/cli.rs +++ b/crates/irtt-app/tests/cli.rs @@ -442,7 +442,7 @@ fn targets_stdin_supersedes_the_active_set_and_stops_on_eof() { let mut stdin = client.0.stdin.take().unwrap(); let mut rows = Vec::new(); for target in ["first", "second"] { - writeln!(stdin, "{target}={}", server.addr).unwrap(); + write!(stdin, "\r\n{target}={}\r\n", server.addr).unwrap(); stdin.flush().unwrap(); let deadline = Instant::now() + Duration::from_secs(5); loop { @@ -485,3 +485,46 @@ fn targets_stdin_supersedes_the_active_set_and_stops_on_eof() { .any(|row| row.contains("\"target\":\"second\"") && row.contains("\"event\":\"session_closed\""))); } + +// Record bounds and decoding errors must terminate the real stdin frontend, +// even when no managed events are available to wake it. +#[cfg(feature = "client")] +#[test] +fn targets_stdin_invalid_records_are_fatal_and_empty_eof_is_normal() { + use std::{ + io::{Read, Write}, + process::Stdio, + }; + for (record, success, diagnostic) in [ + (Vec::new(), true, ""), + (b"\r\n\n".to_vec(), true, ""), + (b"\n=127.0.0.1:9\n".to_vec(), false, "line 2"), + (vec![0xff, b'\n'], false, "line 1"), + (vec![b'x'; 64 * 1024 + 2], false, "line 1"), + ] { + let mut client = ClientProcess( + Command::new(env!("CARGO_BIN_EXE_irtt-client")) + .env_clear() + .args(["--duration", "0", "--targets-stdin"]) + .stdin(Stdio::piped()) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(), + ); + let mut stdin = client.0.stdin.take().unwrap(); + stdin.write_all(&record).unwrap(); + drop(stdin); + let status = client.wait(); + let mut stderr = String::new(); + client + .0 + .stderr + .take() + .unwrap() + .read_to_string(&mut stderr) + .unwrap(); + assert_eq!(status.success(), success, "{stderr}"); + assert!(stderr.contains(diagnostic), "{stderr}"); + } +}