diff --git a/src/format.rs b/src/format.rs index 20908fe..5c54d2d 100644 --- a/src/format.rs +++ b/src/format.rs @@ -17,6 +17,17 @@ pub fn rel_time(d: Duration) -> String { } } +/// Format a byte count in whole binary units for compact status notices. +pub fn bytes(n: usize) -> String { + if n < 1024 { + format!("{n} B") + } else if n < 1024 * 1024 { + format!("{} KiB", n / 1024) + } else { + format!("{} MiB", n / (1024 * 1024)) + } +} + /// Truncate to at most `max` display columns, appending `…` when cut. /// /// Control chars are flattened to spaces so a stray escape/newline from a diff --git a/src/supervisor.rs b/src/supervisor.rs index ac3dd1f..66f6bb1 100644 --- a/src/supervisor.rs +++ b/src/supervisor.rs @@ -177,21 +177,26 @@ impl Supervisor { self.watched = id; } Command::Input { id, bytes } => { - if let Some(t) = self.by_id_mut(id) { - let _ = t.send_input(&bytes); + let refused = self.by_id_mut(id).and_then(|t| t.send_input(&bytes).err()); + if let Some(r) = refused { + self.notice_refused(id, "input", r.len); } } // Paste and scroll land here (not as pre-encoded `Input`) because // their encoding depends on the child's vt100 state, which only // this side of the socket can see. Command::Paste { id, bytes } => { - if let Some(t) = self.by_id_mut(id) { - let _ = t.send_paste(&bytes); + let refused = self.by_id_mut(id).and_then(|t| t.send_paste(&bytes).err()); + if let Some(r) = refused { + self.notice_refused(id, "paste", r.len); } } Command::Mouse { id, kind, col, row } => { - if let Some(t) = self.by_id_mut(id) { - let _ = t.send_mouse(kind, col, row); + let refused = self + .by_id_mut(id) + .and_then(|t| t.send_mouse(kind, col, row).err()); + if let Some(r) = refused { + self.notice_refused(id, "mouse input", r.len); } } Command::Scrollback { id, action } => { @@ -302,6 +307,14 @@ impl Supervisor { // --- internals ------------------------------------------------------------ + /// Report the task and message size for a bounded writer-queue refusal. + fn notice_refused(&mut self, id: u64, what: &str, len: usize) { + self.events.push(Event::Status(format!( + "task {id} is not reading input; dropped {} {what}", + crate::format::bytes(len) + ))); + } + fn index_of(&self, id: u64) -> Option { self.tasks.iter().position(|t| t.id == id) } @@ -706,6 +719,83 @@ mod tests { ); } + /// A blocked PTY write runs off the core thread, so shutdown remains bounded + /// when a child does not read stdin. + #[test] + fn shutdown_survives_a_child_that_never_reads_stdin() { + let mut s = sup(24, 80); + s.set_kill_grace(Duration::from_millis(200)); + s.apply(Command::Spawn { + command: "sleep 300".into(), + cwd: here(), + }); + s.tick(); + let id = match s.drain().first() { + Some(Event::Tasks(v)) => v[0].id, + _ => panic!("expected a Tasks snapshot"), + }; + // Newline-terminated input fills the canonical-mode PTY queue and + // blocks the writer worker while the child is not reading. + s.apply(Command::Paste { + id, + bytes: b"x\n".repeat(1 << 19), + }); + let t0 = Instant::now(); + s.apply(Command::Shutdown); + assert!( + t0.elapsed() < Duration::from_secs(2), + "shutdown blocked behind a PTY write to a non-reading child" + ); + s.tick(); + assert!( + s.drain() + .iter() + .any(|e| matches!(e, Event::Tasks(v) if v.is_empty())) + ); + } + + /// A message that would exceed the writer-queue limit is refused whole, + /// reported with the task ID and size, and does not block the supervisor. + #[test] + fn overfull_writer_queue_refuses_message_with_notice() { + let mut s = sup(24, 80); + s.set_kill_grace(Duration::from_millis(200)); + s.apply(Command::Spawn { + command: "sleep 300".into(), + cwd: here(), + }); + s.tick(); + let id = match s.drain().first() { + Some(Event::Tasks(v)) => v[0].id, + _ => panic!("expected a Tasks snapshot"), + }; + // Newline-terminated input keeps the worker blocked and its admitted + // byte count pending while the child does not read. + let big = b"x\n".repeat(4 << 20); + s.apply(Command::Input { + id, + bytes: big.clone(), + }); + s.apply(Command::Input { + id, + bytes: big.clone(), + }); + s.apply(Command::Input { id, bytes: big }); + let evs = s.drain(); + assert!( + evs.iter().any(|e| matches!(e, Event::Status(m) + if m.contains(&format!("task {id}")) && m.contains("8 MiB"))), + "no refusal notice for the overflowing message; got {evs:?}" + ); + // Shutdown remains bounded after the refusal. + let t0 = Instant::now(); + s.apply(Command::Shutdown); + assert!( + t0.elapsed() < Duration::from_secs(2), + "supervisor wedged after a writer-queue refusal" + ); + } + /// `Shutdown` with a TERM-ignoring job is bounded by the grace, then /// SIGKILLs it: quit can be slowed, never wedged. #[test] diff --git a/src/task.rs b/src/task.rs index 887c863..f3e1b95 100644 --- a/src/task.rs +++ b/src/task.rs @@ -3,6 +3,8 @@ use std::ffi::OsString; use std::io::{self, Read, Write}; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::mpsc::{Sender, channel}; use std::sync::{Arc, Mutex}; use std::thread::{self, JoinHandle}; use std::time::{Duration, Instant}; @@ -18,6 +20,18 @@ use crate::protocol::{Lifecycle, MouseKind, ScrollAction}; /// Number of history rows retained by each task's terminal grid. const SCROLLBACK: usize = 2000; +/// Maximum bytes admitted to one task's writer queue but not yet written to the +/// PTY. This admits one maximum-size paste with headroom while bounding queued +/// input when a child stops reading. +const MAX_PENDING_WRITE: usize = 16 * 1024 * 1024; + +/// A whole-message refusal from the bounded writer queue. +#[derive(Debug)] +pub struct WriteRefused { + /// Size of the refused message, for the client-facing notice. + pub len: usize, +} + /// Map a dependency error (portable-pty returns `anyhow`) into `io::Error` so /// the whole crate speaks stdlib `io::Result` and never grows an `anyhow` dep. fn io_err(e: impl std::fmt::Display) -> io::Error { @@ -151,7 +165,13 @@ pub struct Task { pub cwd: PathBuf, /// Kept for resize (`TIOCSWINSZ`); `try_clone_reader`/`take_writer` borrow it. master: Box, - writer: Box, + /// Sender for the detached PTY writer worker. `None` after `force_kill`. + /// Queuing keeps a blocked PTY write off the core thread. + input_tx: Option>>, + /// Bytes admitted to the writer queue but not yet fully written. Only the + /// core thread admits (single producer), so `queue_write`'s check-then-add + /// cannot over-admit; the worker subtracts after each completed write. + pending_write: Arc, /// Session-leader PID, also used as the process-group ID. pid: Option, /// Shared with the reader thread: it writes (process bytes), the UI reads @@ -294,6 +314,26 @@ impl Task { }) }; + // Drain whole queued messages on a detached worker. The worker is not + // joined because a PTY write can block until the slave side closes. + let (input_tx, input_rx) = channel::>(); + let pending_write = Arc::new(AtomicUsize::new(0)); + { + let pending = Arc::clone(&pending_write); + let mut writer = writer; + thread::spawn(move || { + while let Ok(msg) = input_rx.recv() { + let res = writer.write_all(&msg).and_then(|()| writer.flush()); + pending.fetch_sub(msg.len(), Ordering::Release); + if res.is_err() { + // The slave side closed (child gone): nothing more can + // be delivered. Queued messages drop with the receiver. + break; + } + } + }); + } + // Process-group signalling and `waitid` use the leader PID directly. let pid = child.process_id(); drop(child); @@ -302,7 +342,8 @@ impl Task { command: command.to_string(), cwd: cwd.to_path_buf(), master: pair.master, - writer, + input_tx: Some(input_tx), + pending_write, pid, parser, last_activity, @@ -442,16 +483,37 @@ impl Task { Ok(()) } - pub fn send_input(&mut self, bytes: &[u8]) -> io::Result<()> { - // Input returns the viewport to live before writing to the PTY. - { - let mut p = grid(&self.parser); - if p.screen().scrollback() > 0 { - p.screen_mut().set_scrollback(0); - } + /// Queue `bytes` for the PTY as one message without blocking the caller. + /// Refuse it whole if admission would exceed `MAX_PENDING_WRITE`. + pub fn send_input(&mut self, bytes: &[u8]) -> Result<(), WriteRefused> { + self.snap_live(); + self.queue_write(bytes.to_vec()) + } + + /// Input returns the viewport to live before the bytes are queued. + fn snap_live(&mut self) { + let mut p = grid(&self.parser); + if p.screen().scrollback() > 0 { + p.screen_mut().set_scrollback(0); } - self.writer.write_all(bytes)?; - self.writer.flush() + } + + /// Admit one whole message to the writer queue, or refuse it whole. + fn queue_write(&self, msg: Vec) -> Result<(), WriteRefused> { + // A force-killed task has no writer queue; discard subsequent input. + let Some(tx) = &self.input_tx else { + return Ok(()); + }; + let len = msg.len(); + if self.pending_write.load(Ordering::Acquire) + len > MAX_PENDING_WRITE { + return Err(WriteRefused { len }); + } + self.pending_write.fetch_add(len, Ordering::Release); + if tx.send(msg).is_err() { + // The worker has exited; remove the failed admission from the count. + self.pending_write.fetch_sub(len, Ordering::Release); + } + Ok(()) } /// Move the scrollback viewport, clamped to retained history. @@ -473,17 +535,19 @@ impl Task { } /// Forward a clipboard paste in whichever shape the child negotiated; see - /// [`paste_bytes`]. The grid lock is released before the PTY write: the - /// write can block on a full PTY buffer, and the reader thread needs the - /// lock to drain it. - pub fn send_paste(&mut self, content: &[u8]) -> io::Result<()> { + /// [`paste_bytes`]. Encoded under the grid lock (the DECSET 2004 read stays + /// on this thread), then queued whole: the PTY write itself happens on the + /// writer worker. + pub fn send_paste(&mut self, content: &[u8]) -> Result<(), WriteRefused> { let bracketed = grid(&self.parser).screen().bracketed_paste(); - self.send_input(&paste_bytes(bracketed, content)) + let msg = paste_bytes(bracketed, content); + self.snap_live(); + self.queue_write(msg) } /// Forward one mouse action, routed by the child's own screen state; see /// [`mouse_bytes`]. A child that gets `None` receives nothing at all. - pub fn send_mouse(&mut self, kind: MouseKind, col: u16, row: u16) -> io::Result<()> { + pub fn send_mouse(&mut self, kind: MouseKind, col: u16, row: u16) -> Result<(), WriteRefused> { let bytes = { let p = grid(&self.parser); mouse_bytes(p.screen(), kind, col, row) @@ -541,6 +605,9 @@ impl Task { } self.kill_sent = true; self.handle.take(); // drop the JoinHandle -> detach, never block + // Close the queue without joining a worker that may still be in a PTY + // write. Killing the process group closes the slave side and unblocks it. + self.input_tx.take(); } } @@ -889,6 +956,29 @@ mod tests { t.terminate(); } + /// The per-task writer worker delivers queued messages in FIFO order. + #[test] + fn queued_writes_reach_the_child_in_order() { + let mut t = spawn(10, "cat"); + t.send_input(b"zqfirstqz\n").unwrap(); + t.send_input(b"zqsecondqz\n").unwrap(); + let deadline = Instant::now() + Duration::from_secs(5); + let mut contents = String::new(); + while Instant::now() < deadline { + contents = grid(&t.parser).screen().contents(); + if contents.contains("zqsecondqz") { + break; + } + thread::sleep(Duration::from_millis(10)); + } + let first = contents.find("zqfirstqz").expect("first message never echoed"); + let second = contents + .find("zqsecondqz") + .expect("second message never echoed"); + assert!(first < second, "queued writes reordered: {contents:?}"); + t.terminate(); + } + /// Input hints track child terminal-mode changes. #[test] fn input_hints_track_child_modes() { diff --git a/src/transport.rs b/src/transport.rs index 6869bcf..c197b1e 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -5,6 +5,10 @@ use std::os::unix::net::UnixStream; use std::sync::atomic::AtomicBool; use std::sync::mpsc::{Receiver, Sender, TryRecvError, channel}; use std::thread::{self, JoinHandle}; +use std::time::Duration; + +/// Maximum time allowed for one command-frame write to the daemon. +const SEND_TIMEOUT: Duration = Duration::from_secs(5); use crate::core::{Wake, run_loop}; use crate::frame::{read_frame, write_frame}; @@ -149,7 +153,8 @@ pub struct SocketTransport { write: UnixStream, evt_rx: Receiver, reader: Option>, - /// Set when the reader thread ends on socket EOF: the daemon is gone. + /// Set when the reader thread ends on socket EOF (the daemon is gone), or + /// when a `send` fails (the connection is unrecoverable; see `send`). dead: bool, } @@ -163,6 +168,9 @@ impl SocketTransport { read: UnixStream, wait_tx: Sender<()>, ) -> SocketTransport { + // Keep construction infallible; if this best-effort setup fails, the + // stream retains its existing write-timeout setting. + let _ = write.set_write_timeout(Some(SEND_TIMEOUT)); let (evt_tx, evt_rx) = channel(); let reader = thread::spawn(move || { let mut read = read; @@ -191,9 +199,18 @@ impl SocketTransport { impl Transport for SocketTransport { fn send(&mut self, cmd: Command) { + if self.dead { + // Already unrecoverable; nothing can deliver the command. + return; + } let (kind, payload) = encode_command(&cmd); - // A broken pipe means the daemon is gone; we're tearing down anyway. - let _ = write_frame(&mut self.write, kind, &payload); + if write_frame(&mut self.write, kind, &payload).is_err() { + // A failed frame write may leave a partial frame on the stream. + // Mark the connection dead and close both halves so the reader + // exits and the client can reconnect. + self.dead = true; + let _ = self.write.shutdown(Shutdown::Both); + } } fn poll(&mut self) -> Vec { @@ -254,3 +271,28 @@ impl Transport for LocalTransport { self.sup.apply(Command::Shutdown); } } + +#[cfg(test)] +mod tests { + use super::*; + + /// A failed send marks the transport disconnected and stops its reader. + #[test] + fn failed_send_marks_the_transport_dead() { + let (ours, theirs) = UnixStream::pair().unwrap(); + let write = ours.try_clone().unwrap(); + let (wait_tx, _wait_rx) = channel(); + let mut t = SocketTransport::from_halves(write, ours, wait_tx); + assert!(t.connected()); + + drop(theirs); // the daemon is gone + t.send(Command::Watch { id: None }); + assert!( + !t.connected(), + "a failed send must mark the transport dead" + ); + // The stream was shut down with it, so the reader thread saw EOF and + // exited: joining it cannot hang. + t.reader.take().unwrap().join().unwrap(); + } +}