From 746156bc7c553faa1c19b1e98b921214d9110fe4 Mon Sep 17 00:00:00 2001 From: JJ Liebig Date: Fri, 2 Oct 2026 15:40:19 +0200 Subject: [PATCH] perf(windows): remove the dedicated PTY input forwarding thread (#4850) --- src/pty/actor.rs | 323 ++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 278 insertions(+), 45 deletions(-) diff --git a/src/pty/actor.rs b/src/pty/actor.rs index 76e793ef4..19ce0bf9f 100644 --- a/src/pty/actor.rs +++ b/src/pty/actor.rs @@ -12,7 +12,7 @@ mod windows { use bytes::Bytes; use portable_pty::{MasterPty, PtySize}; - use tokio::sync::mpsc; + use tokio::sync::{mpsc, oneshot}; use tracing::{debug, warn}; pub(crate) struct PtyReadResult { @@ -59,10 +59,16 @@ mod windows { SubmissionPart { bytes: Bytes, deadline: Option, - reply: std_mpsc::Sender>, + reply: oneshot::Sender>, + enter_completion: Option>, }, } + struct InputAcceptance { + accepting: bool, + enter_completion: Option>, + } + enum PtyIoControlCommand { Resize(PtyResizeRequest), Shutdown, @@ -74,7 +80,7 @@ mod windows { control_tx: std_mpsc::Sender, write_tx: std_mpsc::Sender, response_order: Arc>, - accepting: Arc>, + accepting: Arc>, } impl PtyIoActorHandle { @@ -82,10 +88,11 @@ mod windows { &self, bytes: Bytes, ) -> Result<(), mpsc::error::TrySendError> { - if !*self + if !self .accepting .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) + .accepting { return Err(mpsc::error::TrySendError::Closed(bytes)); } @@ -118,7 +125,7 @@ mod windows { .accepting .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); - if !*accepting { + if !accepting.accepting { return Err(std::io::Error::new( std::io::ErrorKind::BrokenPipe, "pty actor closed", @@ -177,9 +184,17 @@ mod windows { } pub(crate) fn shutdown(&self) { - if let Ok(mut accepting) = self.accepting.lock() { - *accepting = false; + let mut accepting = self + .accepting + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + // Pane shutdown terminates the child immediately after this returns. An Enter + // already queued must finish flushing, even if the async task cannot run. + if let Some(completion) = accepting.enter_completion.take() { + let _ = completion.recv(); } + accepting.accepting = false; + drop(accepting); let _ = self.control_tx.send(PtyIoControlCommand::Shutdown); } } @@ -206,7 +221,10 @@ mod windows { let (control_tx, control_rx) = std_mpsc::channel::(); let (write_tx, write_rx) = std_mpsc::channel::(); let response_order = Arc::new(Mutex::new(())); - let accepting = Arc::new(Mutex::new(!initially_quiesced)); + let accepting = Arc::new(Mutex::new(InputAcceptance { + accepting: !initially_quiesced, + enter_completion: None, + })); std::thread::spawn(move || { run_writer(&mut writer, write_rx); @@ -216,9 +234,9 @@ mod windows { { let write_tx = write_tx.clone(); let accepting = Arc::clone(&accepting); - std::thread::spawn(move || { - run_input_forwarder(&mut data_rx, write_tx, accepting); - debug!(pane_id, "windows pty input thread exiting"); + tokio::spawn(async move { + run_input_forwarder(&mut data_rx, write_tx, accepting).await; + debug!(pane_id, "windows pty input task exiting"); }); } @@ -300,6 +318,7 @@ mod windows { bytes, deadline, reply, + enter_completion, } => { let result = if deadline.is_some_and(|deadline| Instant::now() >= deadline) { Err(input_submission_timed_out()) @@ -309,6 +328,9 @@ mod windows { let failed = result .as_ref() .is_err_and(|err| err.kind() != std::io::ErrorKind::TimedOut); + // Wake synchronous shutdown directly; it may be blocking the executor + // worker that would otherwise receive the async acknowledgement. + drop(enter_completion); let _ = reply.send(result); if failed { break; @@ -322,12 +344,12 @@ mod windows { } } - fn run_input_forwarder( + async fn run_input_forwarder( data_rx: &mut mpsc::Receiver, write_tx: std_mpsc::Sender, - accepting: Arc>, + accepting: Arc>, ) { - while let Some(command) = data_rx.blocking_recv() { + while let Some(command) = data_rx.recv().await { match command { PtyIoDataCommand::WriteUserInput(bytes) => { if write_tx.send(PtyIoWriteCommand::Write(bytes)).is_err() { @@ -348,18 +370,16 @@ mod windows { } else { let text_deadline = deadline.and_then(|deadline| deadline.checked_sub(delay)); - write_submission_part(&write_tx, text, text_deadline).and_then(|()| { + async { + write_submission_part(&write_tx, text, text_deadline, None).await?; // A started text write is committed. Finish Enter even if the caller // stops waiting so a timeout cannot leave a partial prompt. - std::thread::sleep(delay); - let accepting = accepting - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if !*accepting { - return Err(pty_actor_closed()); + if !delay.is_zero() { + tokio::time::sleep(delay).await; } - write_submission_part(&write_tx, enter, None) - }) + write_submission_part(&write_tx, enter, None, Some(&accepting)).await + } + .await }; let failed = result .as_ref() @@ -373,22 +393,40 @@ mod windows { } } - fn write_submission_part( + async fn write_submission_part( write_tx: &std_mpsc::Sender, bytes: Bytes, deadline: Option, + accepting: Option<&Arc>>, ) -> std::io::Result<()> { - let (reply, completion) = std_mpsc::channel(); - write_tx - .send(PtyIoWriteCommand::SubmissionPart { - bytes, - deadline, - reply, - }) - .map_err(|_| pty_actor_closed())?; - completion - .recv() - .unwrap_or_else(|_| Err(pty_actor_closed())) + let completion = { + let (reply, completion) = oneshot::channel(); + let mut accepting = accepting.map(|accepting| { + accepting + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + }); + let enter_completion = if let Some(accepting) = accepting.as_mut() { + if !accepting.accepting { + return Err(pty_actor_closed()); + } + let (done, completion) = std_mpsc::channel(); + accepting.enter_completion = Some(completion); + Some(done) + } else { + None + }; + write_tx + .send(PtyIoWriteCommand::SubmissionPart { + bytes, + deadline, + reply, + enter_completion, + }) + .map_err(|_| pty_actor_closed())?; + completion + }; + completion.await.unwrap_or_else(|_| Err(pty_actor_closed())) } fn pty_actor_closed() -> std::io::Error { @@ -410,6 +448,38 @@ mod windows { #[cfg(test)] mod tests { use super::*; + + fn input_acceptance(accepting: bool) -> Arc> { + Arc::new(Mutex::new(InputAcceptance { + accepting, + enter_completion: None, + })) + } + + fn test_runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .unwrap() + } + + async fn next_write_command( + receiver: &std_mpsc::Receiver, + ) -> PtyIoWriteCommand { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + match receiver.try_recv() { + Ok(command) => return command, + Err(std_mpsc::TryRecvError::Empty) => tokio::task::yield_now().await, + Err(err) => panic!("writer disconnected: {err}"), + } + } + }) + .await + .expect("async forwarder queues writer command") + } + struct RecordingWriter { writes: Vec<(Vec, Instant)>, flushes: Vec, @@ -440,7 +510,10 @@ mod windows { fail_after: Option, delay: Duration, deadline: Option, - during_delay: impl FnOnce(&std_mpsc::Sender, &Arc>), + during_delay: impl FnOnce( + &std_mpsc::Sender, + &Arc>, + ), ) -> (RecordingWriter, std::io::Result<()>) { let (flushed_tx, flushed_rx) = std_mpsc::channel(); let mut writer = RecordingWriter { @@ -452,7 +525,7 @@ mod windows { let (data_tx, mut data_rx) = mpsc::channel(2); let (write_tx, write_rx) = std_mpsc::channel(); let (reply_tx, reply_rx) = std_mpsc::channel(); - let accepting = Arc::new(Mutex::new(true)); + let accepting = input_acceptance(true); data_tx .try_send(PtyIoDataCommand::SubmitUserInput { text: Bytes::from_static(b"prompt"), @@ -473,14 +546,15 @@ mod windows { }); let input_write_tx = write_tx.clone(); let input_accepting = Arc::clone(&accepting); - let input_thread = std::thread::spawn(move || { - run_input_forwarder(&mut data_rx, input_write_tx, input_accepting) + let runtime = test_runtime(); + let input_task = runtime.spawn(async move { + run_input_forwarder(&mut data_rx, input_write_tx, input_accepting).await }); flushed_rx.recv().expect("prompt was flushed"); during_delay(&write_tx, &accepting); let result = reply_rx.recv().expect("writer reports submission"); drop(data_tx); - input_thread.join().expect("input thread joins"); + runtime.block_on(input_task).expect("input task joins"); drop(write_tx); (writer_thread.join().expect("writer thread joins"), result) } @@ -515,7 +589,7 @@ mod windows { fn shutdown_during_submission_delay_cancels_enter() { let (writer, result) = run_recorded_submission(None, Duration::from_millis(30), None, |_, accepting| { - *accepting.lock().unwrap() = false; + accepting.lock().unwrap().accepting = false; }); let err = result.expect_err("shutdown cancels enter"); assert_eq!(err.kind(), std::io::ErrorKind::BrokenPipe); @@ -542,7 +616,7 @@ mod windows { let (write_tx, write_rx) = std_mpsc::channel(); let (first_reply_tx, first_reply_rx) = std_mpsc::channel(); let (expired_reply_tx, expired_reply_rx) = std_mpsc::channel(); - let accepting = Arc::new(Mutex::new(true)); + let accepting = input_acceptance(true); data_tx .try_send(PtyIoDataCommand::SubmitUserInput { text: Bytes::from_static(b"first"), @@ -567,15 +641,16 @@ mod windows { writer }); let input_write_tx = write_tx.clone(); - let input_thread = std::thread::spawn(move || { - run_input_forwarder(&mut data_rx, input_write_tx, accepting) + let runtime = test_runtime(); + let input_task = runtime.spawn(async move { + run_input_forwarder(&mut data_rx, input_write_tx, accepting).await }); first_reply_rx.recv().unwrap().unwrap(); let err = expired_reply_rx.recv().unwrap().unwrap_err(); assert_eq!(err.kind(), std::io::ErrorKind::TimedOut); drop(data_tx); - input_thread.join().unwrap(); + runtime.block_on(input_task).unwrap(); drop(write_tx); let writer = writer_thread.join().unwrap(); assert_eq!( @@ -587,6 +662,164 @@ mod windows { vec![b"first".as_slice(), b"\r".as_slice()] ); } + + #[tokio::test] + async fn abandoned_committed_submission_finishes_without_blocking_executor() { + let (data_tx, mut data_rx) = mpsc::channel(2); + let (write_tx, write_rx) = std_mpsc::channel(); + let (reply_tx, reply_rx) = std_mpsc::channel(); + let deadline = Instant::now() + Duration::from_millis(50); + data_tx + .try_send(PtyIoDataCommand::SubmitUserInput { + text: Bytes::from_static(b"committed"), + enter: Bytes::from_static(b"\r"), + delay: Duration::from_millis(20), + deadline: Some(deadline), + reply: reply_tx, + }) + .unwrap(); + data_tx + .try_send(PtyIoDataCommand::WriteUserInput(Bytes::from_static( + b"after", + ))) + .unwrap(); + drop(data_tx); + let forwarder = tokio::spawn(async move { + run_input_forwarder(&mut data_rx, write_tx, input_acceptance(true)).await + }); + + let PtyIoWriteCommand::SubmissionPart { bytes, reply, .. } = + next_write_command(&write_rx).await + else { + panic!("text must be first"); + }; + assert_eq!(bytes, b"committed".as_slice()); + assert!(reply_rx.try_recv().is_err(), "completion waits for flush"); + // A text write can finish after the caller's deadline. Its acknowledgement + // must still be consumed, and Enter must follow even after the caller leaves. + tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await; + drop(reply_rx); + reply.send(Ok(())).unwrap(); + let PtyIoWriteCommand::SubmissionPart { + bytes, + deadline, + reply, + enter_completion, + } = next_write_command(&write_rx).await + else { + panic!("queued user input cannot overtake Enter"); + }; + assert_eq!(bytes, b"\r".as_slice()); + assert!(deadline.is_none(), "committed Enter cannot expire"); + assert!( + write_rx.try_recv().is_err(), + "user input waits for Enter flush" + ); + drop(enter_completion); + reply.send(Ok(())).unwrap(); + forwarder.await.unwrap(); + let PtyIoWriteCommand::Write(bytes) = write_rx.try_recv().unwrap() else { + panic!("ordinary user input follows Enter"); + }; + assert_eq!(bytes, b"after".as_slice()); + } + + #[test] + fn shutdown_waits_for_queued_enter_and_releases_when_writer_finishes_or_disconnects() { + for outcome in ["flushed", "failed", "disconnected"] { + let runtime = test_runtime(); + let accepting = input_acceptance(true); + let (data_tx, _data_rx) = mpsc::channel(1); + let (control_tx, control_rx) = std_mpsc::channel(); + let (write_tx, write_rx) = std_mpsc::channel(); + let handle = PtyIoActorHandle { + data_tx, + control_tx, + write_tx: write_tx.clone(), + response_order: Arc::new(Mutex::new(())), + accepting: Arc::clone(&accepting), + }; + let input_accepting = Arc::clone(&accepting); + let input_task = runtime.spawn(async move { + write_submission_part( + &write_tx, + Bytes::from_static(b"\r"), + None, + Some(&input_accepting), + ) + .await + }); + let command = write_rx.recv_timeout(Duration::from_secs(2)).unwrap(); + let (shutdown_done_tx, shutdown_done_rx) = std_mpsc::channel(); + let shutdown = std::thread::spawn(move || { + handle.shutdown(); + shutdown_done_tx.send(()).unwrap(); + }); + let deadline = Instant::now() + Duration::from_secs(2); + while accepting.try_lock().is_ok() { + assert!(Instant::now() < deadline, "shutdown takes acceptance lock"); + std::thread::yield_now(); + } + assert!( + shutdown_done_rx.try_recv().is_err(), + "shutdown waits for Enter" + ); + assert!( + control_rx.try_recv().is_err(), + "master stays open until Enter finishes" + ); + + let (flushed_tx, _flushed_rx) = std_mpsc::channel(); + let mut writer = RecordingWriter { + writes: vec![], + flushes: vec![], + fail_after: (outcome == "failed").then_some(0), + flushed: flushed_tx, + }; + // Shutdown cannot rely on the async receiver running. Writer completion, + // write failure, and discarded queued commands must all release it. + if outcome == "disconnected" { + drop(command); + } else { + let (writer_tx, writer_rx) = std_mpsc::channel(); + writer_tx.send(command).unwrap(); + drop(writer_tx); + run_writer(&mut writer, writer_rx); + } + shutdown_done_rx + .recv_timeout(Duration::from_secs(2)) + .unwrap(); + shutdown.join().unwrap(); + let result = runtime.block_on(input_task).unwrap(); + assert_eq!(result.is_err(), outcome != "flushed"); + assert_eq!(writer.flushes.len(), usize::from(outcome == "flushed")); + assert!(!accepting.lock().unwrap().accepting); + } + } + + #[tokio::test] + async fn enter_rejects_quiesced_actor_and_failed_enqueue_releases_shutdown_wait() { + let (write_tx, write_rx) = std_mpsc::channel(); + let accepting = input_acceptance(false); + let error = + write_submission_part(&write_tx, Bytes::from_static(b"\r"), None, Some(&accepting)) + .await + .unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::BrokenPipe); + assert!(write_rx.try_recv().is_err()); + accepting.lock().unwrap().accepting = true; + drop(write_rx); + let error = + write_submission_part(&write_tx, Bytes::from_static(b"\r"), None, Some(&accepting)) + .await + .unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::BrokenPipe); + let completion = accepting.lock().unwrap().enter_completion.take().unwrap(); + assert!(matches!( + completion.try_recv(), + Err(std_mpsc::TryRecvError::Disconnected) + )); + } } }