fix: reduce idle ssh cpu without dropping final output (#3728)

* fix: stop polling idle ssh uploads on unix

* fix: preserve ssh output during idle bridge cancellation

* test: exercise active ssh upload cancellation

5 files changed, 366 insertions, 11 deletions

diff --git a/src/client/transport.rs b/src/client/transport.rs
index 4807c31d..2320598b 100644
--- a/src/client/transport.rs
+++ b/src/client/transport.rs
@@ -144,3 +144,83 @@ pub(super) fn write_to_server(
 ) -> io::Result<()> {
     stream.send_client_message(msg)
 }
+
+#[cfg(all(test, unix))]
+mod tests {
+    use super::*;
+    use crate::client::endpoint::EndpointTransport as _;
+    use interprocess::local_socket::traits::Listener as _;
+    use std::io::{Read as _, Write as _};
+    use std::time::Instant;
+
+    #[test]
+    fn upload_cancellation_preserves_pending_endpoint_download() {
+        let path = std::env::temp_dir().join(format!(
+            "herdr-cancel-{}-{}.sock",
+            std::process::id(),
+            std::time::SystemTime::now()
+                .duration_since(std::time::UNIX_EPOCH)
+                .unwrap()
+                .as_nanos()
+        ));
+        let listener = crate::ipc::bind_private_local_listener(&path).unwrap();
+        let client = crate::ipc::connect_local_stream(&path).unwrap();
+        let mut bridge = listener.accept().unwrap();
+        std::fs::remove_file(path).unwrap();
+        drop(listener);
+        let mut reader_stream = client.try_clone().unwrap();
+        let mut writer = endpoint::NativeEndpointTransport::with_lifetime(client, ()).unwrap();
+        let stopped = writer.stop_handle();
+        struct ForwardedInput(std::sync::mpsc::Sender<Vec<u8>>);
+        impl io::Write for ForwardedInput {
+            fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
+                self.0.send(bytes.to_vec()).unwrap();
+                Ok(bytes.len())
+            }
+
+            fn flush(&mut self) -> io::Result<()> {
+                Ok(())
+            }
+        }
+        let (forwarded_tx, forwarded_rx) = std::sync::mpsc::channel();
+        let cancel = crate::remote::bridge_upload_cancellation_for_test(
+            bridge.try_clone().unwrap(),
+            ForwardedInput(forwarded_tx),
+        );
+        let message = ClientMessage::ClientShellFocus { focused: false };
+        let mut expected = Vec::new();
+        protocol::write_message(&mut expected, &message).unwrap();
+        writer.send(&message).unwrap();
+        let mut forwarded = Vec::new();
+        while forwarded.len() < expected.len() {
+            forwarded.extend(forwarded_rx.recv_timeout(Duration::from_secs(3)).unwrap());
+        }
+        assert_eq!(forwarded, expected);
+        cancel();
+
+        // A client write after upload cancellation must not stop the download reader.
+        writer
+            .send(&ClientMessage::ClientShellFocus { focused: true })
+            .unwrap();
+        let flushed = writer.flush(Instant::now() + Duration::from_secs(3));
+        if flushed.is_ok() {
+            let received: ClientMessage =
+                protocol::read_message(&mut bridge, protocol::MAX_FRAME_SIZE).unwrap();
+            assert_eq!(received, ClientMessage::ClientShellFocus { focused: true });
+        }
+        const FINAL: &[u8] = b"pending-download: FINAL OUTPUT\n";
+        bridge.write_all(FINAL).unwrap();
+        drop(bridge);
+        let mut output = Vec::new();
+        EndpointReader {
+            stream: &mut reader_stream,
+            stopped: &stopped,
+        }
+        .read_to_end(&mut output)
+        .unwrap();
+        assert_eq!(output, FINAL);
+        assert!(flushed.is_ok(), "client write failed: {flushed:?}");
+        assert!(!stopped.load(Ordering::Acquire));
+        assert!(writer.take_error().is_none());
+    }
+}
diff --git a/src/platform/mod.rs b/src/platform/mod.rs
index d48ad4fa..335293e8 100644
--- a/src/platform/mod.rs
+++ b/src/platform/mod.rs
@@ -264,7 +264,7 @@ pub(crate) struct RemoteSshConfigPaths {
 #[cfg(unix)]
 mod unix_common;
 #[cfg(unix)]
-pub(crate) use unix_common::{begin_cli_output, end_cli_output};
+pub(crate) use unix_common::{begin_cli_output, end_cli_output, RemoteBridgeWake};
 
 mod client_state;
 pub(crate) use client_state::{create_private_state_file, replace_file, sync_parent_directory};
diff --git a/src/platform/unix_common.rs b/src/platform/unix_common.rs
index 104f6b69..0df215d6 100644
--- a/src/platform/unix_common.rs
+++ b/src/platform/unix_common.rs
@@ -27,6 +27,50 @@ pub(crate) fn wait_client_stream_readable(stream: &crate::ipc::LocalStream) -> s
     Ok(())
 }
 
+pub(crate) struct RemoteBridgeWake {
+    reader: std::os::unix::net::UnixStream,
+    writer: std::os::unix::net::UnixStream,
+}
+
+impl RemoteBridgeWake {
+    pub(crate) fn new() -> std::io::Result<Self> {
+        let (reader, writer) = std::os::unix::net::UnixStream::pair()?;
+        Ok(Self { reader, writer })
+    }
+
+    pub(crate) fn cancel(&self) -> std::io::Result<()> {
+        // EOF stays readable, including when cancellation precedes the wait.
+        self.writer.shutdown(std::net::Shutdown::Write)
+    }
+
+    pub(crate) fn wait(&self, stream: &crate::ipc::LocalStream) -> std::io::Result<()> {
+        use std::os::fd::{AsFd as _, AsRawFd as _};
+        let crate::ipc::LocalStream::UdSocket(stream) = stream;
+        let mut descriptors = [
+            libc::pollfd {
+                fd: stream.as_fd().as_raw_fd(),
+                events: libc::POLLIN,
+                revents: 0,
+            },
+            libc::pollfd {
+                fd: self.reader.as_raw_fd(),
+                events: libc::POLLIN,
+                revents: 0,
+            },
+        ];
+        loop {
+            // SAFETY: both descriptors remain borrowed and the array has two entries.
+            if unsafe { libc::poll(descriptors.as_mut_ptr(), 2, -1) } >= 0 {
+                return Ok(());
+            }
+            let error = std::io::Error::last_os_error();
+            if error.kind() != std::io::ErrorKind::Interrupted {
+                return Err(error);
+            }
+        }
+    }
+}
+
 pub(super) fn read_terminal_grid_size() -> std::io::Result<(u16, u16)> {
     crossterm::terminal::window_size().map(|size| (size.columns, size.rows))
 }
diff --git a/src/platform/windows.rs b/src/platform/windows.rs
index 4deb1864..e58bacae 100644
--- a/src/platform/windows.rs
+++ b/src/platform/windows.rs
@@ -24,6 +24,25 @@ pub(crate) fn classify_child_exit(status: &portable_pty::ExitStatus) -> super::C
     }
 }
 
+pub(crate) struct RemoteBridgeWake;
+
+impl RemoteBridgeWake {
+    pub(crate) fn new() -> std::io::Result<Self> {
+        Ok(Self)
+    }
+
+    pub(crate) fn cancel(&self) -> std::io::Result<()> {
+        // The named-pipe reader checks its cancellation flag between peeks.
+        Ok(())
+    }
+
+    pub(crate) fn wait(&self, _stream: &crate::ipc::LocalStream) -> std::io::Result<()> {
+        // Synchronous named pipes still use peek-before-read polling on Windows.
+        std::thread::sleep(Duration::from_millis(1));
+        Ok(())
+    }
+}
+
 pub(crate) fn wait_client_stream_readable(
     _stream: &crate::ipc::LocalStream,
 ) -> std::io::Result<()> {
diff --git a/src/remote/attach.rs b/src/remote/attach.rs
index 10db88a6..46c0165d 100644
--- a/src/remote/attach.rs
+++ b/src/remote/attach.rs
@@ -1941,6 +1941,63 @@ fn write_managed_ssh_config() -> io::Result<ManagedSshConfig> {
     })
 }
 
+struct BridgeUploadStop {
+    stopped: AtomicBool,
+    wake: crate::platform::RemoteBridgeWake,
+}
+
+impl BridgeUploadStop {
+    fn new() -> io::Result<Self> {
+        Ok(Self {
+            stopped: AtomicBool::new(false),
+            wake: crate::platform::RemoteBridgeWake::new()?,
+        })
+    }
+
+    fn cancel(&self) {
+        if !self.stopped.swap(true, Ordering::AcqRel) {
+            if let Err(error) = self.wake.cancel() {
+                tracing::debug!(%error, "remote bridge read cancellation failed");
+            }
+        }
+    }
+
+    fn is_stopped(&self) -> bool {
+        self.stopped.load(Ordering::Acquire)
+    }
+}
+
+#[cfg(all(test, unix))]
+pub(crate) fn bridge_upload_cancellation_for_test(
+    stream: crate::ipc::LocalStream,
+    mut writer: impl io::Write + Send + 'static,
+) -> impl FnOnce() {
+    stream.set_nonblocking(true).unwrap();
+    let stop = Arc::new(BridgeUploadStop::new().unwrap());
+    let worker_stop = Arc::clone(&stop);
+    let (done_tx, done_rx) = std::sync::mpsc::channel();
+    let worker = thread::spawn(move || {
+        let closed = AtomicBool::new(false);
+        let result = copy_local_stream_to_writer(
+            stream,
+            &mut writer,
+            &worker_stop,
+            &AtomicBool::new(false),
+            &closed,
+        );
+        done_tx
+            .send((result, closed.load(Ordering::Acquire)))
+            .unwrap();
+    });
+    move || {
+        stop.cancel();
+        let (result, closed) = done_rx.recv_timeout(Duration::from_secs(3)).unwrap();
+        worker.join().unwrap();
+        result.unwrap();
+        assert!(!closed, "upload cancellation must not report peer EOF");
+    }
+}
+
 fn bridge_connection(
     stream: crate::ipc::LocalStream,
     target: &str,
@@ -1950,6 +2007,7 @@ fn bridge_connection(
     noninteractive: bool,
     bridge_stop: &Arc<AtomicBool>,
 ) -> io::Result<()> {
+    let upload_stop = Arc::new(BridgeUploadStop::new()?);
     let mut command = Command::new("ssh");
     apply_managed_ssh_options(&mut command, ssh_options);
     if noninteractive {
@@ -1994,7 +2052,6 @@ fn bridge_connection(
     let mut child_to_stream = stream;
 
     let connection_stop = Arc::new(AtomicBool::new(false));
-    let upload_stop = Arc::new(AtomicBool::new(false));
     let upload_failed = Arc::new(AtomicBool::new(false));
     let download_done = Arc::new(AtomicBool::new(false));
     let client_closed = Arc::new(AtomicBool::new(false));
@@ -2025,7 +2082,7 @@ fn bridge_connection(
             &download_bridge_stop,
         );
         download_done_worker.store(true, Ordering::Release);
-        download_upload_stop.store(true, Ordering::Release);
+        download_upload_stop.cancel();
         result
     });
 
@@ -2033,13 +2090,13 @@ fn bridge_connection(
     let (status_result, child_exited) = loop {
         match child.try_wait() {
             Ok(Some(status)) => {
-                upload_stop.store(true, Ordering::Release);
+                upload_stop.cancel();
                 break (Ok(status), true);
             }
             Ok(None) => {}
             Err(err) => {
                 connection_stop.store(true, Ordering::Release);
-                upload_stop.store(true, Ordering::Release);
+                upload_stop.cancel();
                 let _ = child.kill();
                 let _ = child.wait();
                 break (Err(err), false);
@@ -2047,7 +2104,7 @@ fn bridge_connection(
         }
         if bridge_stop.load(Ordering::Acquire) {
             connection_stop.store(true, Ordering::Release);
-            upload_stop.store(true, Ordering::Release);
+            upload_stop.cancel();
             let _ = child.kill();
             break (child.wait(), false);
         }
@@ -2055,7 +2112,7 @@ fn bridge_connection(
             || upload_failed.load(Ordering::Acquire)
             || download_done.load(Ordering::Acquire)
         {
-            upload_stop.store(true, Ordering::Release);
+            upload_stop.cancel();
             let stopped_at = stopped_at.get_or_insert_with(Instant::now);
             if stopped_at.elapsed() >= Duration::from_millis(250) {
                 connection_stop.store(true, Ordering::Release);
@@ -2065,7 +2122,7 @@ fn bridge_connection(
         }
         thread::sleep(BRIDGE_ACCEPT_POLL);
     };
-    upload_stop.store(true, Ordering::Release);
+    upload_stop.cancel();
     if !child_exited {
         connection_stop.store(true, Ordering::Release);
     }
@@ -2144,21 +2201,29 @@ fn copy_reader_to_local_stream<R: io::Read>(
 fn copy_local_stream_to_writer<W: io::Write>(
     mut stream: crate::ipc::LocalStream,
     writer: &mut W,
-    connection_stop: &AtomicBool,
+    connection_stop: &BridgeUploadStop,
     bridge_stop: &AtomicBool,
     client_closed: &AtomicBool,
 ) -> io::Result<u64> {
     let mut buffer = [0_u8; 16 * 1024];
     let mut total = 0;
 
-    while !connection_stop.load(Ordering::Acquire) && !bridge_stop.load(Ordering::Acquire) {
+    while !connection_stop.is_stopped() && !bridge_stop.load(Ordering::Acquire) {
+        #[cfg(all(test, unix))]
+        tests::UPLOAD_READ_ATTEMPTS.with(|attempts| {
+            if let Some(attempts) = attempts.borrow().as_ref() {
+                attempts.fetch_add(1, Ordering::Relaxed);
+            }
+        });
         match crate::ipc::poll_local_stream_read_count(&mut stream, &mut buffer)? {
             crate::ipc::LocalStreamReadCount::Data(read) => {
                 writer.write_all(&buffer[..read])?;
                 writer.flush()?;
                 total += read as u64;
             }
-            crate::ipc::LocalStreamReadCount::Pending => thread::sleep(BRIDGE_IO_POLL),
+            crate::ipc::LocalStreamReadCount::Pending => {
+                connection_stop.wake.wait(&stream)?;
+            }
             crate::ipc::LocalStreamReadCount::Closed => {
                 client_closed.store(true, Ordering::Release);
                 break;
@@ -2247,6 +2312,153 @@ mod tests {
     use super::*;
 
     #[cfg(unix)]
+    thread_local! {
+        pub(super) static UPLOAD_READ_ATTEMPTS: std::cell::RefCell<Option<Arc<std::sync::atomic::AtomicUsize>>> = const { std::cell::RefCell::new(None) };
+    }
+
+    #[cfg(unix)]
+    fn upload_test_streams(name: &str) -> (crate::ipc::LocalStream, crate::ipc::LocalStream) {
+        let socket = local_forward_socket_path(name, "upload-test");
+        let listener = crate::ipc::bind_private_local_listener(&socket).unwrap();
+        let client = crate::ipc::connect_local_stream(&socket).unwrap();
+        let server = listener.accept().unwrap();
+        server.set_nonblocking(true).unwrap();
+        drop(listener);
+        std::fs::remove_file(socket).unwrap();
+        (client, server)
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn bridge_upload_idle_waits_without_repeated_reads_and_cancels() {
+        use std::sync::atomic::AtomicUsize;
+        use std::sync::mpsc;
+
+        let (mut client, stream) = upload_test_streams("idle");
+        let attempts = Arc::new(AtomicUsize::new(0));
+        let worker_attempts = Arc::clone(&attempts);
+        let stop = Arc::new(BridgeUploadStop::new().unwrap());
+        let worker_stop = Arc::clone(&stop);
+        let (done_tx, done_rx) = mpsc::channel();
+        let worker = thread::spawn(move || {
+            UPLOAD_READ_ATTEMPTS.with(|slot| *slot.borrow_mut() = Some(worker_attempts));
+            let mut output = Vec::new();
+            let closed = AtomicBool::new(false);
+            let result = copy_local_stream_to_writer(
+                stream,
+                &mut output,
+                &worker_stop,
+                &AtomicBool::new(false),
+                &closed,
+            );
+            done_tx
+                .send((result, output, closed.load(Ordering::Acquire)))
+                .unwrap();
+        });
+        let deadline = Instant::now() + Duration::from_secs(5);
+        while attempts.load(Ordering::Relaxed) == 0 {
+            assert!(Instant::now() < deadline, "upload worker did not start");
+            thread::sleep(Duration::from_millis(1));
+        }
+        thread::sleep(Duration::from_millis(100));
+        let idle_reads = attempts.load(Ordering::Relaxed);
+        client.write_all(b"pane input").unwrap();
+        let deadline = Instant::now() + Duration::from_secs(5);
+        while attempts.load(Ordering::Relaxed) < idle_reads + 2 {
+            assert!(
+                Instant::now() < deadline,
+                "input did not wake the upload worker"
+            );
+            thread::sleep(Duration::from_millis(1));
+        }
+        thread::sleep(Duration::from_millis(100));
+        let reads_after_input = attempts.load(Ordering::Relaxed);
+        stop.cancel();
+        let (result, output, closed) = done_rx.recv_timeout(Duration::from_secs(2)).unwrap();
+        worker.join().unwrap();
+        assert_eq!(result.unwrap(), 10);
+        assert_eq!(output, b"pane input");
+        assert!(!closed, "cancellation is not a peer disconnect");
+        assert_eq!(idle_reads, 1, "idle forwarding must wait, not retry reads");
+        assert_eq!(
+            reads_after_input, 3,
+            "forwarding must sleep again after input"
+        );
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn bridge_upload_cancel_before_wait_preserves_download() {
+        use std::io::Read as _;
+
+        let (mut client, stream) = upload_test_streams("cancel-before-wait");
+        let mut download = stream.try_clone().unwrap();
+        let stop = BridgeUploadStop::new().unwrap();
+        stop.cancel();
+        stop.cancel();
+        let closed = AtomicBool::new(false);
+        let count = copy_local_stream_to_writer(
+            stream,
+            &mut Vec::new(),
+            &stop,
+            &AtomicBool::new(false),
+            &closed,
+        )
+        .unwrap();
+        assert_eq!(count, 0);
+        assert!(!closed.load(Ordering::Acquire));
+        download.write_all(b"final frame").unwrap();
+        let mut output = [0; 11];
+        client.read_exact(&mut output).unwrap();
+        assert_eq!(&output, b"final frame");
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn bridge_upload_cancel_between_stop_check_and_wait_is_retained() {
+        let (_client, stream) = upload_test_streams("cancel-before-poll");
+        let stop = BridgeUploadStop::new().unwrap();
+        assert!(!stop.is_stopped());
+        stop.cancel();
+        let (done_tx, done_rx) = std::sync::mpsc::channel();
+        let worker = thread::spawn(move || {
+            done_tx.send(stop.wake.wait(&stream)).unwrap();
+        });
+        done_rx
+            .recv_timeout(Duration::from_secs(2))
+            .unwrap()
+            .unwrap();
+        worker.join().unwrap();
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn bridge_upload_drains_input_before_peer_eof() {
+        let (mut client, stream) = upload_test_streams("drain");
+        let payload = vec![b'x'; 1024 * 1024];
+        let expected = payload.clone();
+        let worker = thread::spawn(move || {
+            let stop = BridgeUploadStop::new().unwrap();
+            let mut output = Vec::new();
+            let closed = AtomicBool::new(false);
+            let count = copy_local_stream_to_writer(
+                stream,
+                &mut output,
+                &stop,
+                &AtomicBool::new(false),
+                &closed,
+            )
+            .unwrap();
+            assert!(closed.load(Ordering::Acquire));
+            assert_eq!(count, output.len() as u64);
+            output
+        });
+        client.write_all(&payload).unwrap();
+        drop(client);
+        assert_eq!(worker.join().unwrap(), expected);
+    }
+
+    #[cfg(unix)]
     #[test]
     fn bridge_socket_is_user_only() {
         use std::os::unix::fs::PermissionsExt;