fix: block idle client writer

refs #726

2 files changed, 410 insertions, 75 deletions

diff --git a/src/server/client_transport.rs b/src/server/client_transport.rs
index b7e19bb3..06415d9e 100644
--- a/src/server/client_transport.rs
+++ b/src/server/client_transport.rs
@@ -4,9 +4,11 @@
 //! It converts socket I/O into [`ServerEvent`] values consumed by
 //! `HeadlessServer`.
 
+use std::collections::VecDeque;
 use std::io::{self, Write};
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Arc;
+use std::sync::mpsc::{SendError, TrySendError};
+use std::sync::{Arc, Condvar, Mutex};
 use std::time::Duration;
 
 use interprocess::local_socket::traits::Stream as _;
@@ -44,9 +46,236 @@ const MAX_INPUT_EVENT_BATCH: usize = 4096;
 #[derive(Clone, Debug)]
 pub(crate) struct ClientWriter {
     /// Reliable control messages such as shutdown, notifications, and clipboard writes.
-    pub(crate) control: std::sync::mpsc::Sender<Vec<u8>>,
+    pub(crate) control: ClientControlWriter,
     /// Droppable render messages. Capacity is one so slow clients cannot build lag.
-    pub(crate) render: std::sync::mpsc::SyncSender<Vec<u8>>,
+    pub(crate) render: ClientRenderWriter,
+}
+
+#[cfg(test)]
+impl ClientWriter {
+    pub(crate) fn test_channel(
+        control: std::sync::mpsc::Sender<Vec<u8>>,
+        render: std::sync::mpsc::SyncSender<Vec<u8>>,
+    ) -> Self {
+        Self {
+            control: ClientControlWriter {
+                target: ClientControlTarget::Channel(control),
+            },
+            render: ClientRenderWriter {
+                target: ClientRenderTarget::Channel(render),
+            },
+        }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct ClientControlWriter {
+    target: ClientControlTarget,
+}
+
+#[derive(Debug)]
+enum ClientControlTarget {
+    Queue(Arc<ClientWriterQueue>),
+    #[cfg(test)]
+    Channel(std::sync::mpsc::Sender<Vec<u8>>),
+}
+
+#[derive(Debug)]
+pub(crate) struct ClientRenderWriter {
+    target: ClientRenderTarget,
+}
+
+#[derive(Debug)]
+enum ClientRenderTarget {
+    Queue(Arc<ClientWriterQueue>),
+    #[cfg(test)]
+    Channel(std::sync::mpsc::SyncSender<Vec<u8>>),
+}
+
+impl Clone for ClientControlWriter {
+    fn clone(&self) -> Self {
+        match &self.target {
+            ClientControlTarget::Queue(queue) => {
+                queue.add_sender();
+                Self {
+                    target: ClientControlTarget::Queue(queue.clone()),
+                }
+            }
+            #[cfg(test)]
+            ClientControlTarget::Channel(sender) => Self {
+                target: ClientControlTarget::Channel(sender.clone()),
+            },
+        }
+    }
+}
+
+impl Drop for ClientControlWriter {
+    fn drop(&mut self) {
+        match &self.target {
+            ClientControlTarget::Queue(queue) => queue.remove_sender(),
+            #[cfg(test)]
+            ClientControlTarget::Channel(_) => {}
+        }
+    }
+}
+
+impl ClientControlWriter {
+    fn queue(queue: Arc<ClientWriterQueue>) -> Self {
+        queue.add_sender();
+        Self {
+            target: ClientControlTarget::Queue(queue),
+        }
+    }
+
+    pub(crate) fn send(&self, data: Vec<u8>) -> Result<(), SendError<Vec<u8>>> {
+        match &self.target {
+            ClientControlTarget::Queue(queue) => queue.send_control(data),
+            #[cfg(test)]
+            ClientControlTarget::Channel(sender) => sender.send(data),
+        }
+    }
+}
+
+impl Clone for ClientRenderWriter {
+    fn clone(&self) -> Self {
+        match &self.target {
+            ClientRenderTarget::Queue(queue) => {
+                queue.add_sender();
+                Self {
+                    target: ClientRenderTarget::Queue(queue.clone()),
+                }
+            }
+            #[cfg(test)]
+            ClientRenderTarget::Channel(sender) => Self {
+                target: ClientRenderTarget::Channel(sender.clone()),
+            },
+        }
+    }
+}
+
+impl Drop for ClientRenderWriter {
+    fn drop(&mut self) {
+        match &self.target {
+            ClientRenderTarget::Queue(queue) => queue.remove_sender(),
+            #[cfg(test)]
+            ClientRenderTarget::Channel(_) => {}
+        }
+    }
+}
+
+impl ClientRenderWriter {
+    fn queue(queue: Arc<ClientWriterQueue>) -> Self {
+        queue.add_sender();
+        Self {
+            target: ClientRenderTarget::Queue(queue),
+        }
+    }
+
+    pub(crate) fn try_send(&self, data: Vec<u8>) -> Result<(), TrySendError<Vec<u8>>> {
+        match &self.target {
+            ClientRenderTarget::Queue(queue) => queue.try_send_render(data),
+            #[cfg(test)]
+            ClientRenderTarget::Channel(sender) => sender.try_send(data),
+        }
+    }
+}
+
+#[derive(Debug)]
+struct ClientWriterQueue {
+    state: Mutex<ClientWriterQueueState>,
+    ready: Condvar,
+}
+
+#[derive(Debug, Default)]
+struct ClientWriterQueueState {
+    control: VecDeque<Vec<u8>>,
+    render: Option<Vec<u8>>,
+    senders: usize,
+    writer_alive: bool,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+enum ClientWriteItem {
+    Control(Vec<u8>),
+    Render(Vec<u8>),
+}
+
+impl ClientWriterQueue {
+    fn new() -> Arc<Self> {
+        Arc::new(Self {
+            state: Mutex::new(ClientWriterQueueState {
+                writer_alive: true,
+                ..ClientWriterQueueState::default()
+            }),
+            ready: Condvar::new(),
+        })
+    }
+
+    fn add_sender(&self) {
+        let mut state = self.lock_state();
+        state.senders = state.senders.saturating_add(1);
+    }
+
+    fn remove_sender(&self) {
+        let mut state = self.lock_state();
+        state.senders = state.senders.saturating_sub(1);
+        self.ready.notify_one();
+    }
+
+    fn send_control(&self, data: Vec<u8>) -> Result<(), SendError<Vec<u8>>> {
+        let mut state = self.lock_state();
+        if !state.writer_alive {
+            return Err(SendError(data));
+        }
+        state.control.push_back(data);
+        self.ready.notify_one();
+        Ok(())
+    }
+
+    fn try_send_render(&self, data: Vec<u8>) -> Result<(), TrySendError<Vec<u8>>> {
+        let mut state = self.lock_state();
+        if !state.writer_alive {
+            return Err(TrySendError::Disconnected(data));
+        }
+        if state.render.is_some() {
+            return Err(TrySendError::Full(data));
+        }
+        state.render = Some(data);
+        self.ready.notify_one();
+        Ok(())
+    }
+
+    fn recv(&self) -> Option<ClientWriteItem> {
+        let mut state = self.lock_state();
+        loop {
+            if let Some(data) = state.control.pop_front() {
+                return Some(ClientWriteItem::Control(data));
+            }
+            if let Some(data) = state.render.take() {
+                self.ready.notify_one();
+                return Some(ClientWriteItem::Render(data));
+            }
+            if state.senders == 0 {
+                return None;
+            }
+            state = self
+                .ready
+                .wait(state)
+                .unwrap_or_else(|poisoned| poisoned.into_inner());
+        }
+    }
+
+    fn close_writer(&self) {
+        let mut state = self.lock_state();
+        state.writer_alive = false;
+        self.ready.notify_all();
+    }
+
+    fn lock_state(&self) -> std::sync::MutexGuard<'_, ClientWriterQueueState> {
+        self.state
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
 }
 
 /// Internal event sent from client transport threads to the main event loop.
@@ -306,13 +535,19 @@ pub(crate) fn handle_client_handshake(
     )?;
 
     // Create separate channels for reliable control messages and droppable renders.
-    let (control_tx, control_rx) = std::sync::mpsc::channel::<Vec<u8>>();
-    let (render_tx, render_rx) = std::sync::mpsc::sync_channel::<Vec<u8>>(1);
+    let writer_queue = ClientWriterQueue::new();
     let writer = ClientWriter {
-        control: control_tx,
-        render: render_tx,
+        control: ClientControlWriter::queue(writer_queue.clone()),
+        render: ClientRenderWriter::queue(writer_queue.clone()),
     };
 
+    // Spawn a writer thread that forwards messages from the channels to the stream.
+    let write_stream = stream.try_clone()?;
+    let writer_event_tx = server_event_tx.clone();
+    std::thread::spawn(move || {
+        client_writer_loop(write_stream, client_id, writer_queue, writer_event_tx);
+    });
+
     // Notify the main loop about the new client.
     let _ = server_event_tx.blocking_send(ServerEvent::ClientConnected {
         client_id,
@@ -326,19 +561,6 @@ pub(crate) fn handle_client_handshake(
         writer,
     });
 
-    // Spawn a writer thread that forwards messages from the channels to the stream.
-    let write_stream = stream.try_clone()?;
-    let writer_event_tx = server_event_tx.clone();
-    std::thread::spawn(move || {
-        client_writer_loop(
-            write_stream,
-            client_id,
-            control_rx,
-            render_rx,
-            writer_event_tx,
-        );
-    });
-
     // Enter read loop — read client messages and forward to main loop.
     client_read_loop(stream, client_id, server_event_tx, should_quit)
 }
@@ -347,67 +569,26 @@ pub(crate) fn handle_client_handshake(
 fn client_writer_loop(
     mut stream: LocalStream,
     client_id: u64,
-    control_rx: std::sync::mpsc::Receiver<Vec<u8>>,
-    render_rx: std::sync::mpsc::Receiver<Vec<u8>>,
+    writer_queue: Arc<ClientWriterQueue>,
     server_event_tx: mpsc::Sender<ServerEvent>,
 ) {
-    let mut control_closed = false;
-    let mut render_closed = false;
-
-    loop {
-        match control_rx.try_recv() {
-            Ok(data) => {
+    while let Some(item) = writer_queue.recv() {
+        match item {
+            ClientWriteItem::Control(data) => {
                 if !write_framed_bytes(&mut stream, &data) {
                     break;
                 }
-                continue;
             }
-            Err(std::sync::mpsc::TryRecvError::Empty) => {}
-            Err(std::sync::mpsc::TryRecvError::Disconnected) => control_closed = true,
-        }
-
-        match render_rx.try_recv() {
-            Ok(data) => {
+            ClientWriteItem::Render(data) => {
                 let _ =
                     server_event_tx.blocking_send(ServerEvent::ClientWriterDrained { client_id });
                 if !write_framed_bytes(&mut stream, &data) {
                     break;
                 }
-                continue;
-            }
-            Err(std::sync::mpsc::TryRecvError::Empty) => {}
-            Err(std::sync::mpsc::TryRecvError::Disconnected) => render_closed = true,
-        }
-
-        if control_closed && render_closed {
-            break;
-        }
-
-        if control_closed {
-            match render_rx.recv_timeout(Duration::from_millis(5)) {
-                Ok(data) => {
-                    let _ = server_event_tx
-                        .blocking_send(ServerEvent::ClientWriterDrained { client_id });
-                    if !write_framed_bytes(&mut stream, &data) {
-                        break;
-                    }
-                }
-                Err(std::sync::mpsc::RecvTimeoutError::Timeout) => continue,
-                Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => render_closed = true,
-            }
-            continue;
-        }
-
-        match control_rx.recv_timeout(Duration::from_millis(5)) {
-            Ok(data) => {
-                if !write_framed_bytes(&mut stream, &data) {
-                    break;
-                }
             }
-            Err(std::sync::mpsc::RecvTimeoutError::Timeout) => continue,
-            Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => control_closed = true,
         }
     }
+    writer_queue.close_writer();
     debug!("client writer thread exiting");
 }
 
@@ -600,6 +781,163 @@ mod tests {
         (client, server, TestSocketPath(path))
     }
 
+    fn test_queue_writer() -> (ClientWriter, Arc<ClientWriterQueue>) {
+        let queue = ClientWriterQueue::new();
+        (
+            ClientWriter {
+                control: ClientControlWriter::queue(queue.clone()),
+                render: ClientRenderWriter::queue(queue.clone()),
+            },
+            queue,
+        )
+    }
+
+    fn frame_server_message(message: &ServerMessage) -> Vec<u8> {
+        let mut bytes = Vec::new();
+        protocol::write_message(&mut bytes, message).expect("frame server message");
+        bytes
+    }
+
+    #[test]
+    fn client_writer_queue_keeps_render_slot_bounded() {
+        let (writer, _queue) = test_queue_writer();
+        let first = frame_server_message(&ServerMessage::WindowTitle {
+            title: Some("first".into()),
+        });
+        let second = frame_server_message(&ServerMessage::WindowTitle {
+            title: Some("second".into()),
+        });
+
+        writer.render.try_send(first).expect("first render fits");
+        assert!(matches!(
+            writer.render.try_send(second),
+            Err(TrySendError::Full(_))
+        ));
+    }
+
+    #[test]
+    fn client_writer_prioritizes_control_and_reports_render_drain() {
+        let (mut client_stream, server_stream, _path) = local_stream_pair("client-writer-priority");
+        let (writer, queue) = test_queue_writer();
+        writer
+            .render
+            .try_send(frame_server_message(&ServerMessage::WindowTitle {
+                title: Some("render".into()),
+            }))
+            .expect("queue render");
+        writer
+            .control
+            .send(frame_server_message(&ServerMessage::ReloadSoundConfig))
+            .expect("queue control");
+
+        let (server_event_tx, mut server_event_rx) = mpsc::channel(4);
+        let handle = std::thread::spawn(move || {
+            client_writer_loop(server_stream, 9, queue, server_event_tx);
+        });
+
+        match protocol::read_message(&mut client_stream, MAX_FRAME_SIZE).expect("read control") {
+            ServerMessage::ReloadSoundConfig => {}
+            other => panic!("expected control message first, got {other:?}"),
+        }
+        match protocol::read_message(&mut client_stream, MAX_FRAME_SIZE).expect("read render") {
+            ServerMessage::WindowTitle { title } => assert_eq!(title.as_deref(), Some("render")),
+            other => panic!("expected render message second, got {other:?}"),
+        }
+        match server_event_rx
+            .blocking_recv()
+            .expect("writer drained render slot")
+        {
+            ServerEvent::ClientWriterDrained { client_id } => assert_eq!(client_id, 9),
+            other => panic!("expected writer drained event, got {other:?}"),
+        }
+
+        drop(writer);
+        handle.join().expect("writer exits after senders drop");
+    }
+
+    #[test]
+    fn client_writer_exits_when_all_writer_handles_drop() {
+        let (_client_stream, server_stream, _path) = local_stream_pair("client-writer-drop");
+        let (writer, queue) = test_queue_writer();
+        let (server_event_tx, _server_event_rx) = mpsc::channel(4);
+        let (done_tx, done_rx) = std::sync::mpsc::channel();
+        std::thread::spawn(move || {
+            client_writer_loop(server_stream, 11, queue, server_event_tx);
+            let _ = done_tx.send(());
+        });
+
+        drop(writer);
+        done_rx
+            .recv_timeout(Duration::from_millis(100))
+            .expect("writer exits without polling after senders drop");
+    }
+
+    #[test]
+    fn client_writer_clone_keeps_loop_alive_until_final_drop() {
+        let (mut client_stream, server_stream, _path) =
+            local_stream_pair("client-writer-clone-drop");
+        let (writer, queue) = test_queue_writer();
+        let cloned_writer = writer.clone();
+        let (server_event_tx, _server_event_rx) = mpsc::channel(4);
+        let (done_tx, done_rx) = std::sync::mpsc::channel();
+        std::thread::spawn(move || {
+            client_writer_loop(server_stream, 12, queue, server_event_tx);
+            let _ = done_tx.send(());
+        });
+
+        drop(writer);
+        cloned_writer
+            .control
+            .send(frame_server_message(&ServerMessage::ReloadSoundConfig))
+            .expect("cloned writer still sends after original drops");
+        match protocol::read_message(&mut client_stream, MAX_FRAME_SIZE)
+            .expect("read control from cloned writer")
+        {
+            ServerMessage::ReloadSoundConfig => {}
+            other => panic!("expected cloned control message, got {other:?}"),
+        }
+        assert!(
+            done_rx.recv_timeout(Duration::from_millis(100)).is_err(),
+            "writer exited while cloned handles were still alive"
+        );
+
+        drop(cloned_writer);
+        done_rx
+            .recv_timeout(Duration::from_millis(100))
+            .expect("writer exits after final cloned writer drops");
+    }
+
+    #[test]
+    fn client_writer_closes_queue_after_socket_write_failure() {
+        let (client_stream, server_stream, _path) =
+            local_stream_pair("client-writer-socket-failure");
+        server_stream
+            .set_send_timeout(Some(Duration::from_millis(100)))
+            .expect("set test send timeout");
+        let (writer, queue) = test_queue_writer();
+        let (server_event_tx, _server_event_rx) = mpsc::channel(4);
+        let (done_tx, done_rx) = std::sync::mpsc::channel();
+        std::thread::spawn(move || {
+            client_writer_loop(server_stream, 13, queue, server_event_tx);
+            let _ = done_tx.send(());
+        });
+
+        drop(client_stream);
+        writer
+            .control
+            .send(vec![b'x'; 1024 * 1024])
+            .expect("message is accepted before the writer observes socket failure");
+        done_rx
+            .recv_timeout(Duration::from_secs(1))
+            .expect("writer exits after socket write failure");
+
+        assert!(matches!(writer.control.send(vec![b'y']), Err(SendError(_))));
+        assert!(matches!(
+            writer.render.try_send(vec![b'z']),
+            Err(TrySendError::Disconnected(_))
+        ));
+    }
+
     #[test]
     fn clamp_terminal_size_zero_zero() {
         assert_eq!(
diff --git a/src/server/headless.rs b/src/server/headless.rs
index 35e17efe..0daa2444 100644
--- a/src/server/headless.rs
+++ b/src/server/headless.rs
@@ -4025,10 +4025,7 @@ mod tests {
         let (control_tx, control_rx) = std::sync::mpsc::channel();
         let (render_tx, render_rx) = std::sync::mpsc::sync_channel(1);
         (
-            ClientWriter {
-                control: control_tx,
-                render: render_tx,
-            },
+            ClientWriter::test_channel(control_tx, render_tx),
             control_rx,
             render_rx,
         )
@@ -6308,7 +6305,7 @@ next_tab = ""
             .as_ref()
             .unwrap()
             .render
-            .send(queued)
+            .try_send(queued)
             .expect("pre-fill render queue");
 
         assert!(server.handle_server_event(ServerEvent::ClientInput {
@@ -6430,7 +6427,7 @@ next_tab = ""
             .expect("serialize dummy message");
         client_tx
             .render
-            .send(queued)
+            .try_send(queued)
             .expect("pre-fill render queue");
 
         server.clients.insert(
@@ -6474,7 +6471,7 @@ next_tab = ""
             .expect("serialize dummy message");
         client_tx
             .render
-            .send(queued)
+            .try_send(queued)
             .expect("pre-fill render queue");
 
         server.clients.insert(
@@ -6890,7 +6887,7 @@ next_tab = ""
             .as_ref()
             .unwrap()
             .render
-            .send(queued)
+            .try_send(queued)
             .expect("pre-fill render queue");
         server.app.full_redraw_pending = true;