Fix race between parent and monitor in use_pty

src/exec/use_pty/backchannel.rs

@@ -199,31 +199,31 @@ impl AsFd for ParentBackchannel {
 /// Different messages exchanged between the monitor and the parent process using a [`ParentBackchannel`].
 #[derive(PartialEq, Eq)]
 pub(super) enum MonitorMessage {
-    ExecCommand,
+    Edge,
     Signal(c_int),
 }
 
 impl MonitorMessage {
     const LEN: usize = PREFIX_LEN + MONITOR_DATA_LEN;
-    const EXEC_CMD: Prefix = 0;
+    const EDGE_CMD: Prefix = 0;
     const SIGNAL: Prefix = 1;
 
     fn from_parts(prefix: Prefix, data: MonitorData) -> Self {
         match prefix {
-            Self::EXEC_CMD => Self::ExecCommand,
+            Self::EDGE_CMD => Self::Edge,
             Self::SIGNAL => Self::Signal(data),
             _ => unreachable!(),
         }
     }
 
     fn to_parts(&self) -> (Prefix, MonitorData) {
         let prefix = match self {
-            MonitorMessage::ExecCommand => Self::EXEC_CMD,
+            MonitorMessage::Edge => Self::EDGE_CMD,
             MonitorMessage::Signal(_) => Self::SIGNAL,
         };
 
         let data = match self {
-            MonitorMessage::ExecCommand => 0,
+            MonitorMessage::Edge => 0,
             MonitorMessage::Signal(data) => *data,
         };
 
@@ -234,7 +234,7 @@ impl MonitorMessage {
 impl std::fmt::Debug for MonitorMessage {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::ExecCommand => "ExecCommand".fmt(f),
+            Self::Edge => "Edge".fmt(f),
             &Self::Signal(signal) => write!(f, "Signal({})", signal_fmt(signal)),
         }
     }

src/exec/use_pty/monitor.rs

@@ -75,8 +75,8 @@ pub(super) fn exec_monitor(
         err
     })?;
     // Given that `UnixStream` delivers messages in order it shouldn't be possible to
-    // receive an event different to `ExecCommand` at the beginning.
-    debug_assert_eq!(event, MonitorMessage::ExecCommand);
+    // receive an event different to `Edge` at the beginning.
+    debug_assert_eq!(event, MonitorMessage::Edge);
 
     // FIXME (ogsudo): Some extra config happens here if selinux is available.
 
@@ -170,6 +170,14 @@ pub(super) fn exec_monitor(
         }
     }
 
+    // Wait for the parent to give us red light before shutting down. This avoids missing
+    // output when the monitor exits too quickly.
+    let event = retry_while_interrupted(|| backchannel.recv()).map_err(|err| {
+        dev_warn!("cannot receive red light from parent: {err}");
+        err
+    })?;
+    debug_assert_eq!(event, MonitorMessage::Edge);
+
     // FIXME (ogsudo): The tty is restored here if selinux is available.
 
     // We call `_exit` instead of `exit` to avoid flushing the parent's IO streams by accident.
@@ -309,8 +317,8 @@ impl<'a> MonitorClosure<'a> {
             }
             Ok(event) => {
                 match event {
-                    // We shouldn't receive this event more than once.
-                    MonitorMessage::ExecCommand => unreachable!(),
+                    // We shouldn't receive this event at this point in the protocol
+                    MonitorMessage::Edge => unreachable!(),
                     // Forward signal to the command.
                     MonitorMessage::Signal(signal) => {
                         if let Some(command_pid) = self.command_pid {

src/exec/use_pty/parent.rs

@@ -202,12 +202,10 @@ pub(in crate::exec) fn exec_pty(
     drop(backchannels.monitor);
 
     // Send green light to the monitor after closing the follower.
-    retry_while_interrupted(|| backchannels.parent.send(&MonitorMessage::ExecCommand)).map_err(
-        |err| {
-            dev_error!("cannot send green light to monitor: {err}");
-            err
-        },
-    )?;
+    retry_while_interrupted(|| backchannels.parent.send(&MonitorMessage::Edge)).map_err(|err| {
+        dev_error!("cannot send green light to monitor: {err}");
+        err
+    })?;
 
     let mut closure = ParentClosure::new(
         monitor_pid,
@@ -344,10 +342,19 @@ impl ParentClosure {
     }
 
     fn run(&mut self, registry: EventRegistry<Self>) -> io::Result<ExitReason> {
-        match registry.event_loop(self) {
+        let result = match registry.event_loop(self) {
             StopReason::Break(err) | StopReason::Exit(ParentExit::Backchannel(err)) => Err(err),
             StopReason::Exit(ParentExit::Command(exit_reason)) => Ok(exit_reason),
-        }
+        };
+        // Send red light to the monitor after processing all events
+        retry_while_interrupted(|| self.backchannel.send(&MonitorMessage::Edge)).map_err(
+            |err| {
+                dev_error!("cannot send red light to monitor: {err}");
+                err
+            },
+        )?;
+
+        result
     }
 
     /// Read an event from the backchannel and return the event if it should break the event loop.