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

1 files changed, 80 insertions, 0 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());
+    }
+}