fix(client): allow cancellation while stdout is blocked

This commit is contained in:
Nils Andreas Svee committed 2026-10-28 00:32:52 +07:00
1 parent efa1313e95
commit c8a911297d
2 files changed
+122 -27

No files matched your search

+120 -27
View File
@@ -8,7 +8,7 @@ use std::{
use irtt_client::{ use irtt_client::{
managed::{ managed::{
ManagedClient, ManagedCommandApplyError, ManagedEndReason, ManagedEvent, ManagedClient, ManagedCommandApplyError, ManagedEndReason, ManagedEvent,
ManagedTargetEndReason, ManagedTargetOutcome, TargetInstance, ManagedEventTryRecvError, ManagedTargetEndReason, ManagedTargetOutcome, TargetInstance,
}, },
ClientEvent, ClientEvent,
}; };
@@ -18,8 +18,7 @@ use super::{args::ClientArgs, output::OutputConfig};
use crate::shared::client::{ use crate::shared::client::{
parse_stdin_target_set, parse_stdin_target_set,
session::{ session::{
drain_final_events, peer_close_run_error, request_managed_stop_for_peer_close, peer_close_run_error, request_managed_stop_for_peer_close, request_managed_stop_once,
request_managed_stop_once,
}, },
worker::ManagedWorker, worker::ManagedWorker,
ManagedRunSetup, TargetSpec, STDIN_MAX_DESIRED_TARGETS, STDIN_OUTCOME_HISTORY_LIMIT, ManagedRunSetup, TargetSpec, STDIN_MAX_DESIRED_TARGETS, STDIN_OUTCOME_HISTORY_LIMIT,
@@ -37,6 +36,28 @@ use std::collections::{BTreeMap, BTreeSet};
const MAX_STDIN_RECORD_BYTES: usize = 64 * 1024; const MAX_STDIN_RECORD_BYTES: usize = 64 * 1024;
// The frontend only renders into memory. This thread is the sole stdout owner;
// dropping its JoinHandle deliberately detaches a writer blocked in a pipe.
struct StdoutWriter {
sender: tokio::sync::mpsc::Sender<Vec<u8>>,
}
impl StdoutWriter {
async fn send(&self, buffer: &mut Vec<u8>) -> io::Result<()> {
if !buffer.is_empty() {
// Reserve before taking the buffer: cancellation never discards an
// already accepted chunk, and normal operation preserves ordering.
let permit = self
.sender
.reserve()
.await
.map_err(|_| io::Error::other("stdout writer stopped accepting output"))?;
permit.send(std::mem::take(buffer));
}
Ok(())
}
}
pub async fn run_stream( pub async fn run_stream(
args: ClientArgs, args: ClientArgs,
mut shutdown: tokio::sync::watch::Receiver<bool>, mut shutdown: tokio::sync::watch::Receiver<bool>,
@@ -45,6 +66,53 @@ pub async fn run_stream(
print!("{}", OutputConfig::list_columns()); print!("{}", OutputConfig::list_columns());
return Ok(()); return Ok(());
} }
let (sender, mut receiver) = tokio::sync::mpsc::channel::<Vec<u8>>(64);
let (completed, mut completion) = tokio::sync::oneshot::channel();
let writer = thread::Builder::new()
.name("irtt-stdout".to_owned())
.spawn(move || {
let result = (|| {
let mut stdout = io::LineWriter::new(io::stdout().lock());
while let Some(chunk) = receiver.blocking_recv() {
stdout.write_all(&chunk)?;
}
stdout.flush()
})();
let _ = completed.send(result);
})?;
let output = StdoutWriter { sender };
let run_shutdown = shutdown.clone();
let result = tokio::select! {
biased;
_ = shutdown.wait_for(|requested| *requested) => return Ok(()),
result = &mut completion => {
// Dropping the run future stops and joins its ManagedWorker first.
writer.join().map_err(|_| io::Error::other("stdout writer thread panicked"))?;
result.map_err(|_| io::Error::other("stdout writer completed without a result"))??;
return Err(io::Error::other("stdout writer stopped unexpectedly").into());
}
result = run_stream_output(args, run_shutdown, &output) => result,
};
drop(output);
// The same cancellation guard applies after the managed worker has ended,
// while queued events and summaries are still being written.
let written = tokio::select! {
biased;
_ = shutdown.wait_for(|requested| *requested) => return result,
result = &mut completion => result,
};
writer
.join()
.map_err(|_| io::Error::other("stdout writer thread panicked"))?;
written.map_err(|_| io::Error::other("stdout writer completed without a result"))??;
result
}
async fn run_stream_output(
args: ClientArgs,
mut shutdown: tokio::sync::watch::Receiver<bool>,
output: &StdoutWriter,
) -> Result<(), Box<dyn std::error::Error>> {
let setup = tokio::select! { let setup = tokio::select! {
result = args.prepare() => result.map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?, result = args.prepare() => result.map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?,
_ = shutdown.wait_for(|requested| *requested) => return Ok(()), _ = shutdown.wait_for(|requested| *requested) => return Ok(()),
@@ -63,7 +131,7 @@ pub async fn run_stream(
eprintln!("{warning}"); eprintln!("{warning}");
} }
if setup.stdin_controlled { if setup.stdin_controlled {
return run_stdin_stream(setup, output_config, shutdown, &args).await; return run_stdin_stream(setup, output_config, shutdown, &args, output).await;
} }
let continuous = args.is_continuous(); let continuous = args.is_continuous();
@@ -87,7 +155,7 @@ pub async fn run_stream(
let mut events = handle.subscribe()?; let mut events = handle.subscribe()?;
let mut status = handle.subscribe_status(); let mut status = handle.subscribe_status();
let owner = ManagedWorker::start(task, handle.clone())?; let owner = ManagedWorker::start(task, handle.clone())?;
let mut stdout = io::LineWriter::new(io::stdout().lock()); let mut stdout = Vec::new();
let mut stream_output = StreamOutput { let mut stream_output = StreamOutput {
#[cfg(feature = "stats")] #[cfg(feature = "stats")]
continuous: args.is_continuous(), continuous: args.is_continuous(),
@@ -141,7 +209,10 @@ pub async fn run_stream(
_ = shutdown.changed(), if !interrupted => {} _ = shutdown.changed(), if !interrupted => {}
changed = status.changed() => { if changed.is_err() { break; } } changed = status.changed() => { if changed.is_err() { break; } }
event = events.recv() => match event { event = events.recv() => match event {
Ok(event) => process_event(event, &mut stream_output, #[cfg(feature = "stats")] &mut stats, &mut terminal_targets)?, Ok(event) => {
process_event(event, &mut stream_output, #[cfg(feature = "stats")] &mut stats, &mut terminal_targets)?;
output.send(stream_output.out).await?;
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(count)) => { Err(tokio::sync::broadcast::error::RecvError::Lagged(count)) => {
dropped_events = dropped_events.saturating_add(count); dropped_events = dropped_events.saturating_add(count);
} }
@@ -153,15 +224,24 @@ pub async fn run_stream(
eprintln!("interrupted, closing managed run..."); eprintln!("interrupted, closing managed run...");
} }
let outcome = owner.join().await?; let outcome = owner.join().await?;
drain_final_events(&mut events, &mut dropped_events, |event| { loop {
process_event( match events.try_recv() {
event, Ok(event) => {
&mut stream_output, process_event(
#[cfg(feature = "stats")] event,
&mut stats, &mut stream_output,
&mut terminal_targets, #[cfg(feature = "stats")]
) &mut stats,
})?; &mut terminal_targets,
)?;
output.send(stream_output.out).await?;
}
Err(ManagedEventTryRecvError::Lagged(count)) => {
dropped_events = dropped_events.saturating_add(count);
}
Err(ManagedEventTryRecvError::Empty | ManagedEventTryRecvError::Closed) => break,
}
}
if let Some(warning) = dropped_event_warning(dropped_events) { if let Some(warning) = dropped_event_warning(dropped_events) {
eprintln!("{warning}"); eprintln!("{warning}");
} }
@@ -212,9 +292,10 @@ pub async fn run_stream(
writeln!(stream_output.out, "target: {}", target.label)?; writeln!(stream_output.out, "target: {}", target.label)?;
} }
stream_output.print_summary(target_stats)?; stream_output.print_summary(target_stats)?;
output.send(stream_output.out).await?;
} }
} }
stream_output.out.flush()?; output.send(stream_output.out).await?;
if let Some(error) = terminal_error { if let Some(error) = terminal_error {
return Err(error.into()); return Err(error.into());
} }
@@ -476,6 +557,7 @@ async fn run_stdin_stream(
output_config: OutputConfig, output_config: OutputConfig,
mut shutdown: tokio::sync::watch::Receiver<bool>, mut shutdown: tokio::sync::watch::Receiver<bool>,
args: &ClientArgs, args: &ClientArgs,
output: &StdoutWriter,
) -> Result<(), Box<dyn std::error::Error>> { ) -> Result<(), Box<dyn std::error::Error>> {
if *shutdown.borrow() { if *shutdown.borrow() {
return Ok(()); return Ok(());
@@ -496,7 +578,7 @@ async fn run_stdin_stream(
let (updates, mut stdin) = tokio::sync::watch::channel(StdinUpdate::Targets(Vec::new())); let (updates, mut stdin) = tokio::sync::watch::channel(StdinUpdate::Targets(Vec::new()));
spawn_stdin_target_reader(updates)?; spawn_stdin_target_reader(updates)?;
let mut stdout = io::LineWriter::new(io::stdout().lock()); let mut stdout = Vec::new();
let mut stream_output = StreamOutput { let mut stream_output = StreamOutput {
#[cfg(feature = "stats")] #[cfg(feature = "stats")]
continuous: args.is_continuous(), continuous: args.is_continuous(),
@@ -646,6 +728,7 @@ async fn run_stdin_stream(
event = events.recv() => match event { event = events.recv() => match event {
Ok(event) => { Ok(event) => {
process_stdin_event(event, &mut stream_output, #[cfg(feature = "stats")] &mut stats, &mut terminal_targets)?; process_stdin_event(event, &mut stream_output, #[cfg(feature = "stats")] &mut stats, &mut terminal_targets)?;
output.send(stream_output.out).await?;
#[cfg(feature = "stats")] #[cfg(feature = "stats")]
reconcile_stdin_stats(&mut stats, &status.borrow(), summary_targets.as_ref()); reconcile_stdin_stats(&mut stats, &status.borrow(), summary_targets.as_ref());
} }
@@ -658,15 +741,24 @@ async fn run_stdin_stream(
} }
let outcome = owner.join().await?; let outcome = owner.join().await?;
drain_final_events(&mut events, &mut dropped_events, |event| { loop {
process_stdin_event( match events.try_recv() {
event, Ok(event) => {
&mut stream_output, process_stdin_event(
#[cfg(feature = "stats")] event,
&mut stats, &mut stream_output,
&mut terminal_targets, #[cfg(feature = "stats")]
) &mut stats,
})?; &mut terminal_targets,
)?;
output.send(stream_output.out).await?;
}
Err(ManagedEventTryRecvError::Lagged(count)) => {
dropped_events = dropped_events.saturating_add(count);
}
Err(ManagedEventTryRecvError::Empty | ManagedEventTryRecvError::Closed) => break,
}
}
if let Some(warning) = dropped_event_warning(dropped_events) { if let Some(warning) = dropped_event_warning(dropped_events) {
eprintln!("{warning}"); eprintln!("{warning}");
} }
@@ -703,10 +795,11 @@ async fn run_stdin_stream(
)?; )?;
} }
stream_output.print_summary(stats)?; stream_output.print_summary(stats)?;
output.send(stream_output.out).await?;
} }
} }
} }
stream_output.out.flush()?; output.send(stream_output.out).await?;
match stop { match stop {
Some(StdinStop::Fatal(error)) => Err(error.into()), Some(StdinStop::Fatal(error)) => Err(error.into()),
@@ -1,5 +1,7 @@
#[cfg(feature = "tui")]
use irtt_client::managed::{ManagedEvent, ManagedEventSubscription, ManagedEventTryRecvError}; use irtt_client::managed::{ManagedEvent, ManagedEventSubscription, ManagedEventTryRecvError};
#[cfg(feature = "tui")]
pub(crate) fn drain_final_events<E>( pub(crate) fn drain_final_events<E>(
events: &mut ManagedEventSubscription, events: &mut ManagedEventSubscription,
dropped_events: &mut u64, dropped_events: &mut u64,