diff --git a/crates/lanspread-peer/src/stream_install.rs b/crates/lanspread-peer/src/stream_install.rs index c1fd4ac..08cc962 100644 --- a/crates/lanspread-peer/src/stream_install.rs +++ b/crates/lanspread-peer/src/stream_install.rs @@ -12,7 +12,7 @@ use bytes::Bytes; use crc32fast::Hasher; use futures::{SinkExt, StreamExt}; use lanspread_proto::{Message, Request, StreamInstallFrame}; -use s2n_quic::stream::SendStream; +use s2n_quic::{application, stream::SendStream}; use tokio::{ fs::File, io::{AsyncRead, AsyncReadExt, AsyncWriteExt}, @@ -93,6 +93,7 @@ impl StreamInstallFrameSink { pub async fn send(&self, frame: StreamInstallFrame) -> eyre::Result<()> { tokio::select! { + biased; () = self.cancel_token.cancelled() => { eyre::bail!("streamed install frame send was cancelled"); } @@ -104,6 +105,12 @@ impl StreamInstallFrameSink { } pub trait StreamInstallProvider: Send + Sync { + /// Streams one archive and stops all archive work before returning. + /// + /// Implementations must observe `cancel_token` promptly. In particular, + /// cancellation must terminate and reap any child process before this + /// future resolves so directory-switch draining has a strict quiescence + /// boundary. fn stream_archive<'a>( &'a self, archive: &'a Path, @@ -151,7 +158,7 @@ impl StreamInstallProvider for ExternalUnrarStreamProvider { cancel_token: CancellationToken, ) -> StreamInstallFuture<'a> { Box::pin(async move { - let listing = unrar_listing(&self.program, archive).await?; + let listing = unrar_listing(&self.program, archive, &cancel_token).await?; let archive_name = archive .file_name() .and_then(|name| name.to_str()) @@ -220,14 +227,30 @@ struct RarEntryDraft { crc32: Option, } -async fn unrar_listing(program: &Path, archive: &Path) -> eyre::Result { - let output = Command::new(program) +async fn unrar_listing( + program: &Path, + archive: &Path, + cancel_token: &CancellationToken, +) -> eyre::Result { + let mut child = Command::new(program) .arg("lt") .arg("-cfg-") .arg(archive) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) .kill_on_drop(true) - .output() - .await?; + .spawn()?; + let mut stdout = child + .stdout + .take() + .ok_or_else(|| eyre::eyre!("unrar listing stdout was not captured"))?; + let mut stderr = child + .stderr + .take() + .ok_or_else(|| eyre::eyre!("unrar listing stderr was not captured"))?; + + let output = + capture_unrar_output(&mut child, &mut stdout, &mut stderr, cancel_token, archive).await?; if !output.status.success() { eyre::bail!( "unrar lt failed for {} with status {}: {}", @@ -240,6 +263,80 @@ async fn unrar_listing(program: &Path, archive: &Path) -> eyre::Result, + stderr: Vec, +} + +async fn capture_unrar_output( + child: &mut tokio::process::Child, + stdout: &mut (impl AsyncRead + Unpin), + stderr: &mut (impl AsyncRead + Unpin), + cancel_token: &CancellationToken, + archive: &Path, +) -> eyre::Result { + let mut stdout_bytes = Vec::new(); + let mut stderr_bytes = Vec::new(); + + let capture_result = { + let capture = async { + let (status, _, _) = tokio::try_join!( + child.wait(), + stdout.read_to_end(&mut stdout_bytes), + stderr.read_to_end(&mut stderr_bytes), + )?; + Ok::<_, std::io::Error>(status) + }; + tokio::pin!(capture); + + tokio::select! { + biased; + () = cancel_token.cancelled() => None, + result = &mut capture => Some(result), + } + }; + + match capture_result { + Some(Ok(status)) => Ok(CapturedProcessOutput { + status, + stdout: stdout_bytes, + stderr: stderr_bytes, + }), + Some(Err(capture_error)) => { + if let Err(cleanup_error) = terminate_and_reap_unrar(child, archive).await { + return Err(eyre::eyre!( + "failed to capture unrar listing for {}: {capture_error}; cleanup also failed: {cleanup_error}", + archive.display() + )); + } + Err(eyre::eyre!( + "failed to capture unrar listing for {}: {capture_error}", + archive.display() + )) + } + None => { + terminate_and_reap_unrar(child, archive).await?; + eyre::bail!("streamed archive {} was cancelled", archive.display()); + } + } +} + +async fn terminate_and_reap_unrar( + child: &mut tokio::process::Child, + archive: &Path, +) -> eyre::Result<()> { + let kill_error = child.start_kill().err(); + if let Err(wait_error) = child.wait().await { + eyre::bail!( + "failed to reap unrar listing for {}: {wait_error}; kill error: {kill_error:?}", + archive.display() + ); + } + Ok(()) +} + fn parse_unrar_listing(output: &str) -> eyre::Result { let mut solid = false; let mut entries = Vec::new(); @@ -505,13 +602,16 @@ pub(crate) async fn send_game_install_stream( Ok(archives) => archives, Err(err) => { let message = err.to_string(); - let tx = send_stream_install_error(tx, message.clone()).await; + let tx = + send_stream_install_error_cancellable(tx, message.clone(), game_id, &cancel_token) + .await; return (tx, Err(eyre::eyre!(message))); } }; if archives.is_empty() { let message = format!("no .eti archives found for {game_id}"); - let tx = send_stream_install_error(tx, message.clone()).await; + let tx = send_stream_install_error_cancellable(tx, message.clone(), game_id, &cancel_token) + .await; return (tx, Err(eyre::eyre!(message))); } @@ -519,7 +619,7 @@ pub(crate) async fn send_game_install_stream( let producer_cancel = cancel_token.child_token(); let frame_sink = StreamInstallFrameSink::new(frame_tx, producer_cancel.clone()); let game_id_for_producer = game_id.to_string(); - let producer = tokio::spawn({ + let producer = { let provider = provider.clone(); let producer_cancel = producer_cancel.clone(); async move { @@ -541,34 +641,134 @@ pub(crate) async fn send_game_install_stream( let _ = frame_sink.send(StreamInstallFrame::Complete).await; Ok(()) } - }); + }; + tokio::pin!(producer); let mut framed_tx = FramedWrite::new(tx, LengthDelimitedCodec::new()); - let mut send_result = Ok(()); + let (egress_outcome, producer_result) = { + let egress = forward_stream_install_frames(&mut framed_tx, &mut frame_rx, &producer_cancel); + tokio::pin!(egress); - while let Some(frame) = frame_rx.recv().await { - if let Err(err) = framed_tx.send(frame.encode()).await { - producer_cancel.cancel(); - send_result = Err(eyre::eyre!("failed to send streamed install frame: {err}")); - break; + tokio::select! { + biased; + outcome = &mut egress => (outcome, None), + result = &mut producer => { + let outcome = egress.await; + (outcome, Some(result)) + } } + }; + let egress_result = egress_outcome.into_result(game_id); + + // Once egress stops exceptionally, make the stream unusable before waiting + // for producer cleanup. This prevents queued or partially buffered frames + // from looking like a valid prefix while archive work unwinds. + let mut tx = framed_tx.into_inner(); + if egress_result.is_err() { + producer_cancel.cancel(); + reset_stream_install(&mut tx, game_id); } drop(frame_rx); - let close_result = framed_tx - .close() - .await - .map_err(|err| eyre::eyre!("failed to close streamed install stream: {err}")); - let tx = framed_tx.into_inner(); - let producer_result = match producer.await { - Ok(result) => result, - Err(err) => Err(eyre::eyre!("streamed install producer task failed: {err}")), + let producer_result = match producer_result { + Some(result) => result, + None => producer.await, }; - let result = send_result.and(producer_result).and(close_result); + let result = egress_result.and(producer_result); (tx, result) } +async fn send_stream_install_error_cancellable( + tx: SendStream, + message: String, + game_id: &str, + cancel_token: &CancellationToken, +) -> SendStream { + let (frame_tx, mut frame_rx) = mpsc::channel(1); + frame_tx + .try_send(StreamInstallFrame::Error { message }) + .expect("new one-slot StreamInstall error channel should accept one frame"); + drop(frame_tx); + + let mut framed_tx = FramedWrite::new(tx, LengthDelimitedCodec::new()); + let outcome = forward_stream_install_frames(&mut framed_tx, &mut frame_rx, cancel_token).await; + let mut tx = framed_tx.into_inner(); + if let Err(err) = outcome.into_result(game_id) { + reset_stream_install(&mut tx, game_id); + log::debug!("Failed to send StreamInstall error for {game_id}: {err}"); + } + tx +} + +enum StreamInstallEgressOutcome { + Complete, + Cancelled, + Failed(eyre::Report), +} + +impl StreamInstallEgressOutcome { + fn into_result(self, game_id: &str) -> eyre::Result<()> { + match self { + Self::Complete => Ok(()), + Self::Cancelled => Err(eyre::eyre!("streamed install for {game_id} was cancelled")), + Self::Failed(err) => Err(err), + } + } +} + +async fn forward_stream_install_frames( + framed_tx: &mut FramedWrite, + frame_rx: &mut mpsc::Receiver, + cancel_token: &CancellationToken, +) -> StreamInstallEgressOutcome +where + W: tokio::io::AsyncWrite + Unpin, +{ + loop { + let frame = tokio::select! { + biased; + () = cancel_token.cancelled() => { + return StreamInstallEgressOutcome::Cancelled; + } + frame = frame_rx.recv() => frame, + }; + let Some(frame) = frame else { + break; + }; + + let send_result = tokio::select! { + biased; + () = cancel_token.cancelled() => { + return StreamInstallEgressOutcome::Cancelled; + } + result = framed_tx.send(frame.encode()) => result, + }; + if let Err(err) = send_result { + return StreamInstallEgressOutcome::Failed(eyre::eyre!( + "failed to send streamed install frame: {err}" + )); + } + } + + tokio::select! { + biased; + () = cancel_token.cancelled() => StreamInstallEgressOutcome::Cancelled, + result = framed_tx.close() => match result { + Ok(()) => StreamInstallEgressOutcome::Complete, + Err(err) => StreamInstallEgressOutcome::Failed(eyre::eyre!( + "failed to close streamed install stream: {err}" + )), + }, + } +} + +fn reset_stream_install(tx: &mut SendStream, game_id: &str) { + if let Err(err) = tx.reset(application::Error::UNKNOWN) { + log::debug!("Failed to reset cancelled StreamInstall for {game_id}: {err}"); + } +} + pub(crate) async fn receive_streamed_install( peer_addr: SocketAddr, game_id: &str, @@ -823,9 +1023,272 @@ fn resolve_stream_path(staging_dir: &Path, relative_path: &str) -> eyre::Result< #[cfg(test)] mod tests { + use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + task::{Context, Poll}, + }; + + use tokio::sync::Notify; + use super::*; use crate::test_support::TempDir; + struct PendingWriter { + write_polls: Arc, + first_write: Arc, + } + + struct CancelOnShutdownWriter { + cancel_token: CancellationToken, + } + + struct PendingShutdownWriter { + first_shutdown: Arc, + } + + #[cfg(unix)] + struct FailingReader; + + #[cfg(unix)] + impl tokio::io::AsyncRead for FailingReader { + fn poll_read( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + _buf: &mut tokio::io::ReadBuf<'_>, + ) -> Poll> { + Poll::Ready(Err(std::io::Error::other("synthetic pipe read failure"))) + } + } + + impl tokio::io::AsyncWrite for PendingWriter { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + _buf: &[u8], + ) -> Poll> { + self.write_polls.fetch_add(1, Ordering::SeqCst); + self.first_write.notify_one(); + Poll::Pending + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + impl tokio::io::AsyncWrite for CancelOnShutdownWriter { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Poll::Ready(Ok(buf.len())) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.cancel_token.cancel(); + Poll::Ready(Ok(())) + } + } + + impl tokio::io::AsyncWrite for PendingShutdownWriter { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Poll::Ready(Ok(buf.len())) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.first_shutdown.notify_one(); + Poll::Pending + } + } + + #[tokio::test] + async fn producer_sink_cancellation_wins_over_channel_capacity() { + let (frame_tx, mut frame_rx) = mpsc::channel(1); + let cancel_token = CancellationToken::new(); + cancel_token.cancel(); + let sink = StreamInstallFrameSink::new(frame_tx, cancel_token); + + let err = sink + .send(StreamInstallFrame::Complete) + .await + .expect_err("cancelled producer sink should reject frames"); + + assert!(err.to_string().contains("cancelled")); + assert!(matches!( + frame_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + } + + #[tokio::test] + async fn outbound_cancellation_wins_over_already_queued_frames() { + let (frame_tx, mut frame_rx) = mpsc::channel(2); + frame_tx + .try_send(StreamInstallFrame::Directory { + relative_path: "bin".to_string(), + }) + .expect("first frame should fit"); + frame_tx + .try_send(StreamInstallFrame::Complete) + .expect("second frame should fit"); + + let cancel_token = CancellationToken::new(); + cancel_token.cancel(); + let mut framed_tx = FramedWrite::new(tokio::io::sink(), LengthDelimitedCodec::new()); + + let outcome = + forward_stream_install_frames(&mut framed_tx, &mut frame_rx, &cancel_token).await; + + assert!(matches!(outcome, StreamInstallEgressOutcome::Cancelled)); + assert_eq!(frame_rx.len(), 2, "cancelled egress must not drain frames"); + } + + #[tokio::test] + async fn outbound_cancellation_interrupts_a_blocked_frame_write() { + let (frame_tx, mut frame_rx) = mpsc::channel(1); + frame_tx + .try_send(StreamInstallFrame::Complete) + .expect("frame should fit"); + + let write_polls = Arc::new(AtomicUsize::new(0)); + let first_write = Arc::new(Notify::new()); + let mut framed_tx = FramedWrite::new( + PendingWriter { + write_polls: write_polls.clone(), + first_write: first_write.clone(), + }, + LengthDelimitedCodec::new(), + ); + let cancel_token = CancellationToken::new(); + let cancellation = tokio::spawn({ + let cancel_token = cancel_token.clone(); + async move { + first_write.notified().await; + cancel_token.cancel(); + } + }); + + let outcome = + forward_stream_install_frames(&mut framed_tx, &mut frame_rx, &cancel_token).await; + cancellation + .await + .expect("cancellation helper should finish"); + + assert!(matches!(outcome, StreamInstallEgressOutcome::Cancelled)); + let polls_after_return = write_polls.load(Ordering::SeqCst); + assert!(polls_after_return > 0, "test must reach the blocked write"); + tokio::task::yield_now().await; + assert_eq!( + write_polls.load(Ordering::SeqCst), + polls_after_return, + "no write work may remain after egress returns" + ); + } + + #[tokio::test] + async fn successful_close_is_the_egress_completion_point() { + let (frame_tx, mut frame_rx) = mpsc::channel(1); + drop(frame_tx); + let cancel_token = CancellationToken::new(); + let mut framed_tx = FramedWrite::new( + CancelOnShutdownWriter { + cancel_token: cancel_token.clone(), + }, + LengthDelimitedCodec::new(), + ); + + let outcome = + forward_stream_install_frames(&mut framed_tx, &mut frame_rx, &cancel_token).await; + + assert!(cancel_token.is_cancelled()); + assert!(matches!(outcome, StreamInstallEgressOutcome::Complete)); + } + + #[tokio::test] + async fn outbound_cancellation_interrupts_a_blocked_close() { + let (frame_tx, mut frame_rx) = mpsc::channel(1); + drop(frame_tx); + let first_shutdown = Arc::new(Notify::new()); + let mut framed_tx = FramedWrite::new( + PendingShutdownWriter { + first_shutdown: first_shutdown.clone(), + }, + LengthDelimitedCodec::new(), + ); + let cancel_token = CancellationToken::new(); + let cancellation = tokio::spawn({ + let cancel_token = cancel_token.clone(); + async move { + first_shutdown.notified().await; + cancel_token.cancel(); + } + }); + + let outcome = + forward_stream_install_frames(&mut framed_tx, &mut frame_rx, &cancel_token).await; + cancellation + .await + .expect("cancellation helper should finish"); + + assert!(matches!(outcome, StreamInstallEgressOutcome::Cancelled)); + } + + #[cfg(unix)] + #[tokio::test] + async fn unrar_capture_error_kills_and_reaps_child_before_returning() { + let mut child = Command::new("sleep") + .arg("30") + .kill_on_drop(true) + .spawn() + .expect("sleep process should start"); + assert!( + child + .try_wait() + .expect("initial process status should be readable") + .is_none(), + "test child must still be running before capture" + ); + + let err = capture_unrar_output( + &mut child, + &mut FailingReader, + &mut tokio::io::empty(), + &CancellationToken::new(), + Path::new("broken.eti"), + ) + .await + .expect_err("synthetic pipe failure should fail capture"); + + assert!(err.to_string().contains("synthetic pipe read failure")); + assert!( + child + .try_wait() + .expect("reaped process status should be readable") + .is_some(), + "capture errors must kill and reap the process before returning" + ); + } + #[test] fn stream_paths_stay_inside_staging_dir() { let temp = TempDir::new("lanspread-stream-install-path");