Merge pull request #332 from charliek/feature/plan-024-pty-exit

fix(engine): the reader publishes Exit, so trailing output can't be lost (#255)

Author: charliek

crates/roost-engine/src/pty.rs

@@ -27,13 +27,29 @@ use portable_pty::{native_pty_system, ChildKiller, CommandBuilder, PtySize};
 use tokio::sync::{broadcast, mpsc};
 use tracing::{debug, error, warn};
 
+/// Depth of a tab's output fan-out. A consumer that falls this far
+/// behind gets `RecvError::Lagged` and the skipped bytes are gone for
+/// good — a *second*, independent way a tab's output can be truncated,
+/// unrelated to the exit ordering fixed in #255 (that one is about
+/// ordering, this one about capacity). Deliberately left alone:
+/// closing it means resizing or redesigning the channel, and the only
+/// production consumer (`session.rs`) forwards straight onto an
+/// unbounded mpsc, so it can only lag if the UI's drain stalls. When it
+/// does, `session.rs` reports it as `TabOutput::Error` rather than
+/// silently swallowing it.
 const PTY_OUTPUT_BROADCAST_CAPACITY: usize = 256;
 const PTY_INPUT_CHANNEL_CAPACITY: usize = 64;
 const PTY_OUTPUT_CHUNK_SIZE: usize = 4096;
 /// Grace period after SIGHUP before `close()` escalates to SIGKILL.
 /// Matches the Mac side's 20×10ms teardown window in
 /// `PtySupervisor.swift`.
 const KILL_GRACE: Duration = Duration::from_millis(200);
+/// How long each side of the exit handshake waits on the other before
+/// publishing `Exit` itself (#255). Long enough that the normal path —
+/// the reader hitting EOF within microseconds of the child being
+/// reaped — always wins; short enough that a tab whose reader never
+/// EOFs still reports its exit promptly.
+const EXIT_PUBLISH_GRACE: Duration = Duration::from_millis(250);
 
 /// What a subscriber gets back from `PtySupervisor::subscribe`.
 #[derive(Debug, Clone)]
@@ -44,7 +60,12 @@ pub enum PtyOutputEvent {
     /// per-frame chunks are small and the broadcast clone is cheap
     /// enough at the workloads roost runs).
     Bytes(Vec<u8>),
-    /// PTY child exited with this status.
+    /// PTY child exited with this status. Published by the reader task
+    /// after the last `Bytes` it read, so a consumer that stops here
+    /// has the tab's complete output (#255). The one exception is the
+    /// bounded fallback described on `PtySupervisor::spawn`: a reader
+    /// that never reaches EOF gets `Exit` published out from under it
+    /// on a deadline.
     Exit(i32),
 }
 
@@ -179,6 +200,27 @@ impl PtySupervisor {
     /// into the child as `ROOST_SOCKET` so `roostctl` invoked from
     /// inside the tab dials the right UI.
     ///
+    /// Exit ordering (#255): the reader task publishes `Exit`, after
+    /// the last `Bytes` it read. One producer makes "every byte, then
+    /// the exit" structural instead of a race between the reader and
+    /// the reap task — the shape that used to drop a shell's final
+    /// output. The reap task hands the status over and then waits for
+    /// the reader to finish, but only for `EXIT_PUBLISH_GRACE`: a
+    /// reader can legitimately never reach EOF (a background
+    /// descendant holding the slave fd keeps the master readable
+    /// forever), and an unbounded wait would mean the tab never
+    /// reports its exit. Past the deadline the reap task publishes
+    /// `Exit` itself, and bytes may still follow it — the one
+    /// documented exception to the ordering guarantee. Whichever side
+    /// gets there first, `publish_exit_once` makes it exactly one
+    /// `Exit`.
+    ///
+    /// Session lifetime: the session is installed before the reap task
+    /// starts, and the reap task removes it before it reports the exit
+    /// on either channel. So a session always has a waiter that will
+    /// take it back out, and by the time a consumer sees `Exit` (or
+    /// `TabExited`) the tab is already gone from the map.
+    ///
     /// Errors:
     /// * [`PtyError::DuplicateTab`] — `tab_id` already has a live
     ///   session. Caller must `close()` the prior session first.
@@ -293,10 +335,40 @@ impl PtySupervisor {
 
         let master = pair.master;
 
-        // Reader: blocking read off the master fd, push to broadcast.
+        // The two halves of the exit handshake (#255). `status_*`
+        // carries the reaped status to the reader so it can publish
+        // `Exit` after its final `Bytes`; `reader_alive_*` carries
+        // nothing — the reap task's wait ends on the reader task
+        // dropping its sender, which is precisely "the reader is
+        // done". `exit_published` keeps the two sides to one `Exit`.
+        let (status_tx, status_rx) = std::sync::mpsc::channel::<i32>();
+        let (reader_alive_tx, reader_alive_rx) = std::sync::mpsc::channel::<()>();
+        let exit_published = Arc::new(AtomicBool::new(false));
+
+        // Reader: blocking read off the master fd, push to broadcast,
+        // then publish the exit.
         tokio::task::spawn_blocking({
             let output_tx = output_tx.clone();
-            move || pty_reader_loop(reader_handle, &output_tx, tab_id)
+            let exit_published = exit_published.clone();
+            move || {
+                let _reader_alive = reader_alive_tx;
+                pty_reader_loop(reader_handle, &output_tx, tab_id);
+                // EOF: everything the PTY produced is on the channel,
+                // so `Exit` published from here can only follow it.
+                // The status normally lands within microseconds (the
+                // reap task's `waitpid` is already blocked when the
+                // child dies); if it doesn't, the reap task publishes
+                // once it does, still after this EOF.
+                match status_rx.recv_timeout(EXIT_PUBLISH_GRACE) {
+                    Ok(status) => {
+                        publish_exit_once(&output_tx, &exit_published, status);
+                    }
+                    Err(_) => debug!(
+                        tab_id,
+                        "pty reader finished before the child was reaped; reap task publishes Exit"
+                    ),
+                }
+            }
         });
 
         // Writer + resizer: a single ordered loop over the unified
@@ -323,13 +395,59 @@ impl PtySupervisor {
             debug!(tab_id, "pty input loop ended");
         });
 
-        // Wait for the child to exit; forward the exit status onto
-        // the output channel AND the lifecycle channel so both
-        // per-tab consumers and the workspace converge.
         let output_tx_exit = output_tx.clone();
         let lifecycle_tx = self.lifecycle.clone();
         let sessions_for_reap = self.sessions.clone();
         let reaped_for_wait = reaped.clone();
+        let exit_published_for_wait = exit_published.clone();
+
+        let session = Session {
+            cmd_tx,
+            output_tx,
+            killer: Mutex::new(killer),
+            pid,
+            reaped,
+            initial_rx: Some(initial_rx),
+        };
+        // Promote the slot from pending → sessions atomically, BEFORE
+        // the reap task exists (see below).
+        //
+        // If `close(tab_id)` ran while we were building the PTY it
+        // removed our entry from `pending` as a cancellation signal.
+        // Detect that here and hand the session back instead of
+        // installing it; the caller-visible teardown happens after the
+        // reap task is started, so the child is still reaped.
+        let cancelled = {
+            let mut sessions = self.sessions.lock().unwrap();
+            let mut pending = self.pending.lock().unwrap();
+            if pending.remove(&tab_id) {
+                sessions.insert(tab_id, session);
+                None
+            } else {
+                Some(session)
+            }
+        };
+        // Either branch consumed the pending entry (ours, or the one
+        // `close()` already took), so the guard has nothing left to do.
+        slot.armed = false;
+
+        // Wait for the child to exit; hand the status to the reader
+        // task (which publishes it onto the output channel) and send
+        // it on the lifecycle channel so both per-tab consumers and
+        // the workspace converge.
+        //
+        // Started only now that the promotion has run, because this
+        // task's identity-checked removal is the ONLY thing that ever
+        // takes the session back out. Starting it earlier meant a
+        // child that exited during the promotion window got reaped
+        // first: the removal found no session and removed nothing,
+        // then the promotion installed a session whose child was
+        // already dead — `has()` kept answering yes and `write()` kept
+        // accepting input for a PTY nobody was reading. Ordering it
+        // after the insert makes "a reaped child leaves no session"
+        // structural. It also means no `Exit` can be published before
+        // the session is reachable: the reader only publishes once
+        // this task hands it a status.
         tokio::task::spawn_blocking(move || {
             let status = match child.wait() {
                 Ok(status) => status.exit_code() as i32,
@@ -339,66 +457,58 @@ impl PtySupervisor {
                 }
             };
             // Mark reaped first so a concurrent `close()` SIGKILL
-            // watchdog stands down, then publish exit and drop the
-            // dead session so later writes get `NotFound` instead of
-            // silently succeeding against a closed PTY.
+            // watchdog stands down, then drop the dead session so
+            // later writes get `NotFound` instead of silently
+            // succeeding against a closed PTY — and only then tell
+            // anyone the child exited. Removing ahead of both the
+            // status handoff and `TabExited` means the tab is already
+            // unreachable by the time either channel reports the exit,
+            // so a consumer reacting to `Exit` can never find a live
+            // session for a dead child.
             reaped_for_wait.store(true, Ordering::SeqCst);
-            let _ = output_tx_exit.send(PtyOutputEvent::Exit(status));
+            {
+                // Only remove the session if THIS waiter still owns it.
+                // `close()` frees the slot synchronously, so the same
+                // tab_id can be re-spawned before a stale waiter fires;
+                // matching the per-spawn `reaped` identity prevents
+                // evicting a newer live session (#80). Scoped so the
+                // deadline wait below never holds the sessions lock.
+                let mut sessions = sessions_for_reap.lock().unwrap();
+                let owns = sessions
+                    .get(&tab_id)
+                    .map(|s| Arc::ptr_eq(&s.reaped, &reaped_for_wait))
+                    .unwrap_or(false);
+                if owns {
+                    sessions.remove(&tab_id);
+                }
+            }
+            let _ = status_tx.send(status);
             let _ = lifecycle_tx.send(SupervisorEvent::TabExited { tab_id, status });
-            // Only remove the session if THIS waiter still owns it.
-            // `close()` frees the slot synchronously, so the same
-            // tab_id can be re-spawned before a stale waiter fires;
-            // matching the per-spawn `reaped` identity prevents
-            // evicting a newer live session (#80).
-            let mut sessions = sessions_for_reap.lock().unwrap();
-            let owns = sessions
-                .get(&tab_id)
-                .map(|s| Arc::ptr_eq(&s.reaped, &reaped_for_wait))
-                .unwrap_or(false);
-            if owns {
-                sessions.remove(&tab_id);
+            // Backstop for a reader that never reaches EOF (#255).
+            // Ends as soon as the reader task drops its sender —
+            // by then it has published `Exit` and this is a no-op —
+            // or on the deadline, when publishing here is the only
+            // way the tab ever reports its exit.
+            let _ = reader_alive_rx.recv_timeout(EXIT_PUBLISH_GRACE);
+            if publish_exit_once(&output_tx_exit, &exit_published_for_wait, status) {
+                debug!(
+                    tab_id,
+                    "pty reader had not finished; published Exit on the deadline path"
+                );
             }
         });
 
-        let session = Session {
-            cmd_tx,
-            output_tx,
-            killer: Mutex::new(killer),
-            pid,
-            reaped,
-            initial_rx: Some(initial_rx),
-        };
-        // Promote the slot from pending → sessions atomically.
-        // If `close(tab_id)` ran while we were building the PTY it
-        // removed our entry from `pending` as a cancellation
-        // signal. Detect that here, kill the freshly-spawned
-        // child, and don't insert into `sessions`. The killer was
-        // moved into the wait task already, so we reach for the
-        // copy we stashed in `session` below — actually the
-        // session struct already holds the killer, so we tear it
-        // back down via `terminate_child` (SIGHUP→SIGKILL) and drop
-        // `session` (which drops the input/resize channels, the
-        // writer task exits, and the wait task reaps once the
-        // signal lands).
-        {
-            let mut sessions = self.sessions.lock().unwrap();
-            let mut pending = self.pending.lock().unwrap();
-            if !pending.remove(&tab_id) {
-                // Cancelled by close(). Kill the child rather than
-                // returning a usable receiver.
-                drop(pending);
-                drop(sessions);
-                terminate_child(&session.killer, session.pid, session.reaped.clone(), tab_id);
-                drop(session);
-                // SlotGuard is no longer needed — pending was
-                // already cleared by close(); we already cleaned
-                // up the child.
-                slot.armed = false;
-                return Err(PtyError::Cancelled(tab_id).into());
-            }
-            sessions.insert(tab_id, session);
+        if let Some(session) = cancelled {
+            // Cancelled by close(). Kill the child rather than
+            // returning a usable receiver: `terminate_child` sends
+            // SIGHUP (SIGKILL on the watchdog), and the reap task
+            // started above reaps whatever the signal lands on.
+            // Dropping `session` drops the input/resize channels, so
+            // the writer task exits too.
+            terminate_child(&session.killer, session.pid, session.reaped.clone(), tab_id);
+            drop(session);
+            return Err(PtyError::Cancelled(tab_id).into());
         }
-        slot.armed = false;
 
         Ok(early_rx)
     }
@@ -801,6 +911,26 @@ fn build_command(
     cmd
 }
 
+/// Publish a tab's `Exit` event, at most once per spawn. Both the
+/// reader task (the normal path) and the reap task's deadline backstop
+/// call this; the compare-exchange decides which one gets to send, so
+/// a consumer never sees two exits for one child. Returns whether this
+/// call was the one that published.
+fn publish_exit_once(
+    output_tx: &broadcast::Sender<PtyOutputEvent>,
+    published: &AtomicBool,
+    status: i32,
+) -> bool {
+    if published
+        .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
+        .is_err()
+    {
+        return false;
+    }
+    let _ = output_tx.send(PtyOutputEvent::Exit(status));
+    true
+}
+
 fn pty_reader_loop(
     mut reader: Box<dyn Read + Send>,
     output_tx: &broadcast::Sender<PtyOutputEvent>,

crates/roost-engine/src/session.rs

@@ -89,13 +89,34 @@ impl TabSession {
                             break;
                         }
                     }
+                    // Stopping here is safe on the normal path: the
+                    // supervisor's reader task publishes `Exit` after
+                    // the last bytes it read, so nothing is left
+                    // behind (#255).
+                    //
+                    // The exception is `pty.rs`'s bounded deadline
+                    // fallback. A reader that never reaches EOF — a
+                    // background descendant holding the slave fd keeps
+                    // the master readable forever — has `Exit`
+                    // published out from under it after
+                    // `EXIT_PUBLISH_GRACE`, and the bytes it reads
+                    // afterwards are dropped by the `break` below.
+                    // That is the deliberate trade: a tab that never
+                    // reports its exit would never auto-close.
                     Ok(PtyOutputEvent::Exit(status)) => {
                         let _ = output_tx.send(TabOutput::Exit {
                             status,
                             reason: String::new(),
                         });
                         break;
                     }
+                    // The other way a tab's output can be truncated,
+                    // independent of the #255 ordering fix: this drain
+                    // fell far enough behind that the broadcast
+                    // dropped `n` messages. Out of scope there —
+                    // fixing it means resizing or redesigning the
+                    // channel (see `PTY_OUTPUT_BROADCAST_CAPACITY`).
+                    // Surfaced rather than swallowed.
                     Err(RecvError::Lagged(n)) => {
                         let _ = output_tx.send(TabOutput::Error(format!(
                             "broadcast lagged: dropped {n} message(s)"

crates/roost-engine/tests/pty_env_terminfo.rs

@@ -24,13 +24,13 @@ async fn inherited_terminfo_is_stripped_from_child_env() {
         .spawn(11, "/tmp", &["/usr/bin/env".into()], 80, 24, &socket)
         .expect("spawn");
 
-    // Same budget-bounded drain rationale as pty_smoke's
-    // collect_until_closed: content assertions below are what prove the
-    // capture wasn't truncated.
+    // Same drain rationale as pty_smoke's collect_until_exit: `Exit`
+    // arrives after the reader's last bytes (#255), and the content
+    // assertions below are what prove the capture wasn't truncated.
     let deadline = Instant::now() + Duration::from_secs(5);
     let mut collected = Vec::new();
     let mut exit_status = None;
-    while Instant::now() < deadline {
+    while Instant::now() < deadline && exit_status.is_none() {
         match output.try_recv() {
             Ok(PtyOutputEvent::Bytes(bytes)) => collected.extend_from_slice(&bytes),
             Ok(PtyOutputEvent::Exit(status)) => exit_status = Some(status),