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/run.rs b/crates/irtt-app/src/cmd/client/run.rs index f6ac03c..5fc15b0 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, + BlockingManagedClient, ManagedCommandApplyError, ManagedEndReason, ManagedEvent, + ManagedEventSubscription, ManagedEventTryRecvError, ManagedStatus, ManagedTargetConfig, + ManagedTargetEndReason, ManagedTargetOutcome, TargetInstance, }, ClientEvent, }; @@ -21,10 +20,10 @@ use super::{ }; 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 +39,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 +61,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,7 +69,7 @@ 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( @@ -79,6 +77,7 @@ pub fn run_stream( setup.managed_targets(), )?; let handle = owner.handle(); + let mut status = handle.subscribe_status(); let mut stdout = io::LineWriter::new(io::stdout().lock()); let mut stream_output = StreamOutput { config: output_config, @@ -102,9 +101,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 +112,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_async().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 +153,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 +195,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 +276,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 +291,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 +392,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,27 +432,12 @@ 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( @@ -641,8 +445,9 @@ fn run_stdin_stream( setup.managed_targets(), )?; let handle = owner.handle(); - let mailbox = Arc::new(StdinMailbox::default()); - spawn_stdin_target_reader(Arc::clone(&mailbox))?; + let mut status = handle.subscribe_status(); + 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 +459,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_async().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 +631,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/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}"); + } +} diff --git a/crates/irtt-client/src/managed/blocking.rs b/crates/irtt-client/src/managed/blocking.rs index 192cfe8..3d170c3 100644 --- a/crates/irtt-client/src/managed/blocking.rs +++ b/crates/irtt-client/src/managed/blocking.rs @@ -4,7 +4,10 @@ use std::{ }; use thiserror::Error; -use tokio::runtime::{Builder, Runtime}; +use tokio::{ + runtime::{Builder, Runtime}, + sync::oneshot, +}; use super::{ ManagedClient, ManagedClientConfig, ManagedClientHandle, ManagedClientTask, ManagedConfigError, @@ -83,16 +86,19 @@ pub enum BlockingManagedJoinError { WorkerPanicked, } -/// Blocking completion owner for one unified managed Tokio task. +/// Dedicated-thread completion owner for one unified managed Tokio task. /// /// This owner runs exactly one [`ManagedClient`] task on one dedicated OS thread -/// with its own current-thread Tokio runtime. Dropping a live owner requests a -/// graceful stop and synchronously joins that worker, so drop may block while +/// with its own current-thread Tokio runtime. Synchronous callers use [`Self::join`]; +/// async frontends use [`Self::join_async`] without sharing the measurement executor. +/// Dropping a live owner requests graceful stop and synchronously joins that +/// worker, so drop may block while /// bounded managed cleanup completes. #[must_use = "dropping the blocking managed owner requests stop and joins its worker"] pub struct BlockingManagedClient { handle: ManagedClientHandle, worker: Option>, + completion: oneshot::Receiver<()>, } impl BlockingManagedClient { @@ -131,11 +137,13 @@ impl BlockingManagedClient { .build() .map_err(|source| BlockingManagedStartError::Runtime { source })?; let runtime = WorkerRuntime::new(runtime); + let (completed, completion) = oneshot::channel(); let worker = thread::Builder::new() .name("irtt-managed".to_owned()) .spawn(move || { let outcome = runtime.block_on(task); runtime.shutdown(); + let _ = completed.send(()); outcome }) .map_err(|source| BlockingManagedStartError::ThreadSpawn { source })?; @@ -143,6 +151,7 @@ impl BlockingManagedClient { Ok(Self { handle, worker: Some(worker), + completion, }) } @@ -156,6 +165,15 @@ impl BlockingManagedClient { self.join_worker() } + /// Await worker completion, then join its OS thread and obtain its outcome. + /// + /// The task and runtime remain isolated from the caller's executor. Cancelling + /// this future drops the owner, requesting stop and joining as usual. + pub async fn join_async(mut self) -> Result { + let _ = (&mut self.completion).await; + self.join_worker() + } + fn join_worker(&mut self) -> Result { self.worker .take() diff --git a/crates/irtt-client/tests/managed_control.rs b/crates/irtt-client/tests/managed_control.rs index 37f2511..4002bc5 100644 --- a/crates/irtt-client/tests/managed_control.rs +++ b/crates/irtt-client/tests/managed_control.rs @@ -197,3 +197,25 @@ fn concurrent_submission_and_stop_or_seal_never_lose_an_accepted_receipt() { .expect("every accepted concurrent submission must receive a terminal disposition"); }); } + +// Joining an isolated worker must leave the caller's executor free to stop it. +#[test] +fn async_worker_join_allows_concurrent_stop_and_preserves_final_status() { + let owner = BlockingManagedClient::start(config(), vec![]).unwrap(); + let handle = owner.handle(); + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + let joined = owner.join_async(); + let stop = async { handle.stop().await }; + let (outcome, ()) = + tokio::time::timeout(Duration::from_secs(2), async { tokio::join!(joined, stop) }) + .await + .expect("async join must allow stop on the caller's executor"); + let outcome = outcome.unwrap(); + assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested); + assert_eq!(handle.status().final_outcome.as_deref(), Some(&outcome)); + }); +}