fix: batch queued endpoint input frames (#4004)

refs #3833

Co-authored-by: akbash-bot <300245827+akbash-bot@users.noreply.github.com>

1 files changed, 185 insertions, 28 deletions

diff --git a/src/client/endpoint/writer.rs b/src/client/endpoint/writer.rs
index c4e94923..56c472cd 100644
--- a/src/client/endpoint/writer.rs
+++ b/src/client/endpoint/writer.rs
@@ -9,20 +9,30 @@ use super::EndpointTransport;
 use crate::ipc::LocalStream;
 use crate::protocol::ClientMessage;
 
-const MAX_QUEUED_MESSAGES: usize = 256;
+const MAX_QUEUED_BATCHES: usize = 256;
+const MAX_BATCH_BYTES: usize = 64 * 1024;
 const MAX_QUEUED_BYTES: usize = 2 * crate::protocol::MAX_GRAPHICS_FRAME_SIZE;
 const WRITE_TIMEOUT: Duration = Duration::from_secs(5);
 const IO_POLL_INTERVAL: Duration = Duration::from_millis(2);
 
+#[derive(Default)]
+struct FrameBatch {
+    frames: Vec<Vec<u8>>,
+    bytes: usize,
+}
+
 enum WriterCommand {
-    Frame(Vec<u8>),
+    Frames(Arc<Mutex<FrameBatch>>),
     Flush(mpsc::Sender<()>),
 }
 
-/// The UI only enqueues complete frames. A worker owns partial writes, cancellation, and the
-/// bridge lifetime, so neither socket backpressure nor bridge teardown can block other endpoints.
+/// The UI batches complete frames until the worker claims them, so a short burst of tiny input
+/// frames does not exhaust command slots. Frames retain their individual write boundaries.
+/// A worker owns partial writes, cancellation, and the bridge lifetime; socket backpressure and
+/// bridge teardown never block other endpoints.
 pub(crate) struct NativeEndpointTransport {
     sender: mpsc::SyncSender<WriterCommand>,
+    pending_batch: Option<Arc<Mutex<FrameBatch>>>,
     queued_bytes: Arc<AtomicUsize>,
     stopped: Arc<AtomicBool>,
     error: Arc<Mutex<Option<io::Error>>>,
@@ -34,7 +44,7 @@ impl NativeEndpointTransport {
         lifetime: impl Send + 'static,
     ) -> io::Result<Self> {
         stream.set_nonblocking(true)?;
-        let (sender, receiver) = mpsc::sync_channel::<WriterCommand>(MAX_QUEUED_MESSAGES);
+        let (sender, receiver) = mpsc::sync_channel::<WriterCommand>(MAX_QUEUED_BATCHES);
         let queued_bytes = Arc::new(AtomicUsize::new(0));
         let stopped = Arc::new(AtomicBool::new(false));
         let error = Arc::new(Mutex::new(None));
@@ -49,15 +59,14 @@ impl NativeEndpointTransport {
                     if worker_stop.load(Ordering::Acquire) {
                         break;
                     }
-                    let frame = match command {
-                        WriterCommand::Frame(frame) => frame,
+                    let batch = match command {
+                        WriterCommand::Frames(batch) => batch,
                         WriterCommand::Flush(done) => {
                             let _ = done.send(());
                             continue;
                         }
                     };
-                    let result = write_frame(&mut stream, &frame, &worker_stop);
-                    worker_bytes.fetch_sub(frame.len(), Ordering::AcqRel);
+                    let result = write_batch(&mut stream, &batch, &worker_stop, &worker_bytes);
                     if let Err(error) = result {
                         if let Ok(mut slot) = worker_error.lock() {
                             *slot = Some(error);
@@ -69,12 +78,43 @@ impl NativeEndpointTransport {
             })?;
         Ok(Self {
             sender,
+            pending_batch: None,
             queued_bytes,
             stopped,
             error,
         })
     }
 
+    fn enqueue_frame(&mut self, frame: Vec<u8>) -> io::Result<()> {
+        if let Some(batch) = &self.pending_batch {
+            let mut batch = batch
+                .lock()
+                .map_err(|_| io::Error::other("endpoint batch lock poisoned"))?;
+            // An empty batch has already been claimed by the worker. Never append to it.
+            if !batch.frames.is_empty()
+                && frame.len() <= MAX_BATCH_BYTES.saturating_sub(batch.bytes)
+            {
+                batch.bytes += frame.len();
+                batch.frames.push(frame);
+                return Ok(());
+            }
+        }
+        let batch = Arc::new(Mutex::new(FrameBatch {
+            bytes: frame.len(),
+            frames: vec![frame],
+        }));
+        self.sender
+            .try_send(WriterCommand::Frames(batch.clone()))
+            .map_err(|error| match error {
+                mpsc::TrySendError::Full(_) => queue_full(),
+                mpsc::TrySendError::Disconnected(_) => {
+                    io::Error::new(io::ErrorKind::BrokenPipe, "endpoint writer stopped")
+                }
+            })?;
+        self.pending_batch = Some(batch);
+        Ok(())
+    }
+
     pub(crate) fn stop_handle(&self) -> Arc<AtomicBool> {
         self.stopped.clone()
     }
@@ -103,17 +143,9 @@ impl EndpointTransport for NativeEndpointTransport {
         {
             return Err(queue_full());
         }
-        self.sender
-            .try_send(WriterCommand::Frame(frame))
-            .map_err(|error| {
-                self.queued_bytes.fetch_sub(len, Ordering::AcqRel);
-                match error {
-                    mpsc::TrySendError::Full(_) => queue_full(),
-                    mpsc::TrySendError::Disconnected(_) => {
-                        io::Error::new(io::ErrorKind::BrokenPipe, "endpoint writer stopped")
-                    }
-                }
-            })
+        self.enqueue_frame(frame).inspect_err(|_| {
+            self.queued_bytes.fetch_sub(len, Ordering::AcqRel);
+        })
     }
 
     fn disconnect(&mut self) {
@@ -121,6 +153,8 @@ impl EndpointTransport for NativeEndpointTransport {
     }
 
     fn flush(&mut self, deadline: Instant) -> io::Result<()> {
+        // Later frames must stay after the flush command, even if its wait times out.
+        self.pending_batch = None;
         let (done, completion) = mpsc::channel();
         self.sender
             .try_send(WriterCommand::Flush(done))
@@ -157,6 +191,29 @@ fn queue_full() -> io::Error {
     )
 }
 
+fn write_batch(
+    writer: &mut impl io::Write,
+    batch: &Mutex<FrameBatch>,
+    stopped: &AtomicBool,
+    queued_bytes: &AtomicUsize,
+) -> io::Result<()> {
+    // Claim the frames before doing any I/O. The producer never waits for socket progress.
+    let batch = std::mem::take(
+        &mut *batch
+            .lock()
+            .map_err(|_| io::Error::other("endpoint batch lock poisoned"))?,
+    );
+    for frame in batch.frames {
+        if stopped.load(Ordering::Acquire) {
+            break;
+        }
+        let result = write_frame(writer, &frame, stopped);
+        queued_bytes.fetch_sub(frame.len(), Ordering::AcqRel);
+        result?;
+    }
+    Ok(())
+}
+
 fn write_frame(
     writer: &mut impl io::Write,
     mut frame: &[u8],
@@ -404,20 +461,120 @@ mod tests {
         worker.join().unwrap();
     }
 
+    fn queued_transport(
+        capacity: usize,
+    ) -> (NativeEndpointTransport, mpsc::Receiver<WriterCommand>) {
+        let (sender, receiver) = mpsc::sync_channel(capacity);
+        (
+            NativeEndpointTransport {
+                sender,
+                pending_batch: None,
+                queued_bytes: Arc::new(AtomicUsize::new(0)),
+                stopped: Arc::new(AtomicBool::new(false)),
+                error: Arc::new(Mutex::new(None)),
+            },
+            receiver,
+        )
+    }
+
     #[test]
-    fn a_full_queue_is_a_connection_failure_not_silent_input_loss() {
-        let (sender, _receiver) = mpsc::sync_channel(1);
-        let mut transport = NativeEndpointTransport {
-            sender,
-            queued_bytes: Arc::new(AtomicUsize::new(0)),
-            stopped: Arc::new(AtomicBool::new(false)),
-            error: Arc::new(Mutex::new(None)),
-        };
+    fn stdin_burst_is_queued_in_order_without_worker_progress() {
+        let (mut transport, receiver) = queued_transport(MAX_QUEUED_BATCHES);
+        let input = (0..128)
+            .map(|index| format!("{index:04}: ordered input burst\n"))
+            .collect::<String>();
+        let mut framer = crate::raw_input::RawInputByteFramer::for_host_input();
+        let mut expected = Vec::new();
+        for data in framer.push(input.as_bytes()) {
+            let message = ClientMessage::Input { data };
+            crate::protocol::write_message(&mut expected, &message).unwrap();
+            transport
+                .send(&message)
+                .expect("a small stdin burst must fit");
+        }
+        let queued_bytes = transport.queued_bytes.clone();
+        assert_eq!(queued_bytes.load(Ordering::Acquire), expected.len());
+        let mut received = Vec::new();
+        for command in receiver.try_iter() {
+            let WriterCommand::Frames(batch) = command else {
+                panic!("unexpected flush");
+            };
+            assert!(batch.lock().unwrap().bytes <= MAX_BATCH_BYTES);
+            write_batch(&mut received, &batch, &transport.stopped, &queued_bytes).unwrap();
+        }
+        assert_eq!(received, expected);
+        assert_eq!(queued_bytes.load(Ordering::Acquire), 0);
+
+        // The producer still holds the last claimed batch; new input must get a new command.
         transport.send(&ClientMessage::Detach).unwrap();
+        let WriterCommand::Frames(batch) = receiver.try_recv().unwrap() else {
+            panic!("expected a new batch after the worker claimed the previous one");
+        };
+        write_batch(&mut received, &batch, &transport.stopped, &queued_bytes).unwrap();
+        crate::protocol::write_message(&mut expected, &ClientMessage::Detach).unwrap();
+        assert_eq!(received, expected);
+        assert_eq!(queued_bytes.load(Ordering::Acquire), 0);
+    }
+
+    #[test]
+    fn later_frames_cannot_join_a_batch_before_a_flush() {
+        let (mut transport, receiver) = queued_transport(3);
+        let first = ClientMessage::ClientShellFocus { focused: true };
+        let last = ClientMessage::Detach;
+        transport.send(&first).unwrap();
+        assert_eq!(
+            transport.flush(Instant::now()).unwrap_err().kind(),
+            io::ErrorKind::TimedOut
+        );
+        transport.send(&last).unwrap();
+
+        let mut received = Vec::new();
+        let WriterCommand::Frames(batch) = receiver.try_recv().unwrap() else {
+            panic!("expected first batch");
+        };
+        write_batch(
+            &mut received,
+            &batch,
+            &transport.stopped,
+            &transport.queued_bytes,
+        )
+        .unwrap();
+        let mut expected = Vec::new();
+        crate::protocol::write_message(&mut expected, &first).unwrap();
+        assert_eq!(received, expected);
+        assert!(matches!(
+            receiver.try_recv().unwrap(),
+            WriterCommand::Flush(_)
+        ));
+        let WriterCommand::Frames(batch) = receiver.try_recv().unwrap() else {
+            panic!("expected a separate batch after flush");
+        };
+        write_batch(
+            &mut received,
+            &batch,
+            &transport.stopped,
+            &transport.queued_bytes,
+        )
+        .unwrap();
+        crate::protocol::write_message(&mut expected, &last).unwrap();
+        assert_eq!(received, expected);
+        assert_eq!(transport.queued_bytes.load(Ordering::Acquire), 0);
+    }
+
+    #[test]
+    fn a_full_queue_is_a_connection_failure_not_silent_input_loss() {
+        let (mut transport, _receiver) = queued_transport(1);
+        transport
+            .send(&ClientMessage::Input {
+                data: vec![b'x'; MAX_BATCH_BYTES],
+            })
+            .unwrap();
+        let queued = transport.queued_bytes.load(Ordering::Acquire);
         assert_eq!(
             transport.send(&ClientMessage::Detach).unwrap_err().kind(),
             io::ErrorKind::ConnectionAborted
         );
+        assert_eq!(transport.queued_bytes.load(Ordering::Acquire), queued);
         transport
             .queued_bytes
             .store(MAX_QUEUED_BYTES, Ordering::Release);