fix(nesinit): one thing reaps, and the workload stops alone

Three problems in the shutdown and reaping paths, all of them found in review.

Reaping and waiting cannot be two mechanisms. `waitpid(-1, ...)` collects any
child, so the reaper and a caller waiting on its own child race for the same
status, and whichever loses gets nothing — losing the workload's exit, which is
the one thing this component exists to report. The reaper is now the only
waiter and hands each exit to whoever asked for that pid. Registering interest
holds the same lock the delivery takes, so a child that exits before its caller
is registered is still delivered rather than dropped; there is a test that
fails without that.

Killing the workload killed everything. `kill(-1, SIGKILL)` is every process
init may signal, so a workload that overstayed its grace period took the
guest's services with it, before the ordered stop the shutdown promises them
had even started. It signals the one pid now.

The shutdown had no workload to stop. It built a fresh handle with no pid, so
the graceful stop was a no-op and the workload only died in the sweep that
follows — which is exactly the order this was written to avoid. The handle the
session used is now the handle the shutdown uses, and waiting for the workload
waits for that pid rather than for any child to leave.

apps/nesinit/src/main.rs

@@ -6,7 +6,7 @@
 
 use std::time::Duration;
 
-use nesinit::reap;
+use nesinit::reap::{self, Waiters};
 use nesinit::session::{self, Outcome};
 use nesinit::shutdown::{self, Machine};
 use nesinit::workload::{Process, Workload};
@@ -40,26 +40,34 @@ fn main() -> anyhow::Result<()> {
         tracing::warn!(pid, "not PID 1: the kernel will reparent orphans elsewhere");
     }
 
+    let waiters = Waiters::new();
+    let mut workload = Process::new(waiters.clone());
+
     let runtime = tokio::runtime::Builder::new_multi_thread()
         .enable_all()
         .build()?;
-    let outcome = runtime.block_on(guest());
+    let outcome = runtime.block_on(guest(&waiters, &mut workload));
     match &outcome {
         Ok(outcome) => tracing::info!(?outcome, "the session ended"),
         Err(error) => tracing::error!(%error, "the session failed"),
     }
 
+    // Before anything below waits on a pid: the reaper runs on this runtime's
+    // threads, and two things calling `wait` is what the registry exists to
+    // prevent.
+    drop(runtime);
+
     // Reached however the session ended, including an error: an init that
     // returns leaves the guest running with nothing in it.
-    let mut machine = Guest::new();
+    let mut machine = Guest { workload };
     shutdown::ordered(&mut machine, GRACE);
     unreachable!("power_off does not return");
 }
 
-async fn guest() -> anyhow::Result<Outcome> {
+async fn guest(waiters: &Waiters, workload: &mut Process) -> anyhow::Result<Outcome> {
     // Reaping runs for as long as the guest does. A workload that leaks
     // orphans leaks them while it is running, not when it stops.
-    tokio::spawn(reaper());
+    tokio::spawn(reaper(waiters.clone()));
 
     let address = VsockAddr::new(VMADDR_CID_HOST, CONTROL_PORT);
     // Dialled once, with no retry: the far end is listening before this
@@ -67,9 +75,8 @@ async fn guest() -> anyhow::Result<Outcome> {
     // waiting will not fix.
     let channel = VsockStream::connect(address).await?;
 
-    let mut workload = Process::new();
     let outcome = tokio::select! {
-        outcome = session::run(channel, &mut workload) => outcome?,
+        outcome = session::run(channel, workload) => outcome?,
         signal = asked_to_stop() => {
             signal?;
             tracing::info!("asked to stop");
@@ -79,8 +86,9 @@ async fn guest() -> anyhow::Result<Outcome> {
     Ok(outcome)
 }
 
-/// Drain exited children whenever the kernel says there are some.
+/// Drain exited children whenever the kernel says there are some, and hand
-async fn reaper() {
+/// each exit to whoever is waiting for it.
+async fn reaper(waiters: Waiters) {
     let mut children = match signal(SignalKind::child()) {
         Ok(children) => children,
         Err(error) => {
@@ -91,7 +99,10 @@ async fn reaper() {
     loop {
         children.recv().await;
         for (pid, exit) in reap::reap_exited() {
-            tracing::debug!(pid, ?exit, "reaped");
+            if !waiters.deliver(pid, exit) {
+                // An orphan nothing asked about, which is most of them.
+                tracing::debug!(pid, ?exit, "reaped");
+            }
         }
     }
 }
@@ -112,14 +123,6 @@ struct Guest {
     workload: Process,
 }
 
-impl Guest {
-    fn new() -> Self {
-        Self {
-            workload: Process::new(),
-        }
-    }
-}
-
 impl Machine for Guest {
     fn signal_workload(&mut self) {
         self.workload.signal_stop();
@@ -127,18 +130,23 @@ impl Machine for Guest {
 
     fn await_workload(&mut self, grace: Duration) -> bool {
         // The session already reported the exit if there was one; this is the
-        // window for a workload that was asked to stop on the way down.
+        // window for a workload that was asked to stop on the way down. It
-        wait_for_quiet(grace)
+        // waits for that pid and no other, or the first service to leave would
+        // look like the workload leaving.
+        self.workload.await_exit(grace)
     }
 
     fn kill_workload(&mut self) {
-        // SAFETY: two integers, and the kernel refuses to signal init itself.
+        // The workload alone. Everything else in the guest is still expected
-        unsafe { libc::kill(-1, libc::SIGKILL) };
+        // to get the ordered stop below, and a kill to every process here
+        // would take the services with it.
+        self.workload.signal(libc::SIGKILL);
+        self.workload.await_exit(Duration::from_secs(1));
     }
 
     fn signal_rest(&mut self, grace: Duration) {
         // -1 is every process this one may signal, which as PID 1 is all of
-        // them but itself.
+        // them but itself. The workload has already stopped by here.
         unsafe { libc::kill(-1, libc::SIGTERM) };
         wait_for_quiet(grace);
     }

apps/nesinit/src/reap.rs

@@ -5,9 +5,12 @@
 // for. Zombies hold a pid and a slot in the process table; a workload that
 // leaks them in a long session eventually cannot fork.
 
+use std::collections::HashMap;
 use std::io;
+use std::sync::{Arc, Mutex};
 
 use nesprotocol::lifecycle::Exit;
+use tokio::sync::oneshot;
 
 /// Reap every child that has already exited, without blocking on any that
 /// have not.
@@ -70,3 +73,54 @@ pub fn become_subreaper() -> io::Result<()> {
 pub fn refuse_oom_kill() -> io::Result<()> {
     std::fs::write("/proc/self/oom_score_adj", "-1000\n")
 }
+
+/// The children whose exit somebody is waiting for.
+///
+/// Reaping and waiting cannot be two mechanisms. `waitpid(-1, ...)` collects
+/// any child, so a reaper running beside a caller that waits on its own child
+/// will sometimes collect that child first and the caller's wait then fails
+/// with no status — the exit is gone, and an exit is the one thing this
+/// component exists to report. So the reaper is the only waiter, and this is
+/// how it hands an exit to whoever asked for one.
+#[derive(Clone, Default)]
+pub struct Waiters(Arc<Mutex<HashMap<i32, oneshot::Sender<Exit>>>>);
+
+impl Waiters {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Start something and be waiting for it before it can be reaped.
+    ///
+    /// The lock is held across the spawn, and that is the whole point: a child
+    /// that exits immediately is reaped by a thread that has to take the same
+    /// lock to deliver the exit, so it waits until the slot exists rather than
+    /// finding none and dropping it.
+    pub fn watch<F>(&self, spawn: F) -> io::Result<oneshot::Receiver<Exit>>
+    where
+        F: FnOnce() -> io::Result<i32>,
+    {
+        let mut waiting = self.lock();
+        let pid = spawn()?;
+        let (sender, receiver) = oneshot::channel();
+        waiting.insert(pid, sender);
+        Ok(receiver)
+    }
+
+    /// Hand an exit to whoever is waiting for that pid. `false` if nobody is,
+    /// which is the common case: most of what an init reaps is an orphan
+    /// nothing asked about.
+    pub fn deliver(&self, pid: i32, exit: Exit) -> bool {
+        let Some(sender) = self.lock().remove(&pid) else {
+            return false;
+        };
+        // An error here means the caller stopped waiting, which is its right.
+        sender.send(exit).is_ok()
+    }
+
+    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<i32, oneshot::Sender<Exit>>> {
+        self.0
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+}

apps/nesinit/src/workload.rs

@@ -12,6 +12,10 @@ use std::pin::Pin;
 
 use nesprotocol::lifecycle::{Exec, Exit, Mount};
 
+use std::os::unix::process::CommandExt;
+
+use crate::reap::Waiters;
+
 /// Why something could not be done, in the words the operating system used.
 ///
 /// The reason is passed through verbatim on purpose: `EACCES` and a path can
@@ -50,18 +54,50 @@ pub trait Workload {
 
 /// The workload as a local process.
 pub struct Process {
+    waiters: Waiters,
     pid: Option<i32>,
 }
 
 impl Process {
-    pub fn new() -> Self {
+    /// Started through the reaper's registry, because the reaper is the only
-        Self { pid: None }
+    /// thing in this component that may call `wait`.
+    pub fn new(waiters: Waiters) -> Self {
+        Self { waiters, pid: None }
     }
-}
 
-impl Default for Process {
+    /// Send a signal to the workload and nothing else.
-    fn default() -> Self {
+    ///
-        Self::new()
+    /// Nothing here ever signals the process group or every process: the order
+    /// a shutdown promises is only true if the workload can be stopped alone.
+    pub fn signal(&self, signal: libc::c_int) {
+        let Some(pid) = self.pid else { return };
+        // SAFETY: two integers, and it cannot touch this process's memory. A
+        // pid that has already gone fails with ESRCH, which is exactly the
+        // idempotence the callers of this rely on.
+        unsafe { libc::kill(pid, signal) };
+    }
+
+    /// Wait up to `grace` for the workload to leave, without a runtime.
+    ///
+    /// Used on the way down, after the reaper has stopped: it polls for the
+    /// one pid rather than for any child, so a slow service cannot be mistaken
+    /// for the workload still running.
+    pub fn await_exit(&self, grace: std::time::Duration) -> bool {
+        let Some(pid) = self.pid else { return true };
+        let deadline = std::time::Instant::now() + grace;
+        loop {
+            let mut status: libc::c_int = 0;
+            // SAFETY: waitpid writes only into `status`.
+            let seen = unsafe { libc::waitpid(pid, &mut status, libc::WNOHANG) };
+            // -1 is ECHILD: already reaped, which is also gone.
+            if seen == pid || seen == -1 {
+                return true;
+            }
+            if std::time::Instant::now() >= deadline {
+                return false;
+            }
+            std::thread::sleep(std::time::Duration::from_millis(50));
+        }
     }
 }
 
@@ -84,7 +120,10 @@ impl Workload for Process {
             return Err(Failure::new("the command is empty"));
         };
 
-        let mut command = tokio::process::Command::new(program);
+        // The standard library's process rather than the runtime's: the
+        // runtime reaps the children it spawns, and in this component reaping
+        // belongs to one place. See `reap::Waiters`.
+        let mut command = std::process::Command::new(program);
         command.args(args);
         // Cleared rather than inherited: init's environment is the kernel's
         // and says nothing a workload should read.
@@ -112,34 +151,27 @@ impl Workload for Process {
             });
         }
 
-        let mut child = command.spawn().map_err(|e| Failure::new(e.to_string()))?;
+        let mut pid = None;
-        self.pid = child.id().map(|pid| pid as i32);
+        let exited = self
-        let waiter = async move {
+            .waiters
-            let status = child.wait().await?;
+            .watch(|| {
-            Ok(exit_of(status))
+                let child = command.spawn()?;
-        };
+                let started = child.id() as i32;
-        Ok(Box::pin(waiter))
+                pid = Some(started);
+                Ok(started)
+            })
+            .map_err(|error| Failure::new(error.to_string()))?;
+        self.pid = pid;
+
+        Ok(Box::pin(async move {
+            exited
+                .await
+                .map_err(|_| io::Error::other("the workload's exit was not delivered"))
+        }))
     }
 
     fn signal_stop(&mut self) {
-        if let Some(pid) = self.pid {
+        self.signal(libc::SIGTERM);
-            // SAFETY: kill takes two integers and cannot touch this process's
-            // memory. A pid that has already exited fails with ESRCH, which is
-            // exactly the idempotence this method promises.
-            unsafe { libc::kill(pid, libc::SIGTERM) };
-        }
-    }
-}
-
-/// Read a finished process's status as an exit.
-fn exit_of(status: std::process::ExitStatus) -> Exit {
-    use std::os::unix::process::ExitStatusExt;
-    match (status.code(), status.signal()) {
-        (Some(code), _) => Exit::code(code),
-        (None, Some(signal)) => Exit::signal(signal),
-        // Neither: nothing on this platform produces it, and inventing a zero
-        // would report a clean run.
-        (None, None) => Exit::signal(0),
     }
 }
 

apps/nesinit/tests/reaping.rs

@@ -3,10 +3,11 @@
 // Its own test binary on purpose: the reaper waits on any child, so it would
 // collect processes another test in the same binary was waiting for.
 
-use std::sync::{Mutex, MutexGuard};
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::{Arc, Mutex, MutexGuard};
 use std::time::{Duration, Instant};
 
-use nesinit::reap::{become_subreaper, reap_exited};
+use nesinit::reap::{Waiters, become_subreaper, reap_exited};
 
 /// Reaping is process-wide: `waitpid(-1, ...)` collects any child, so two of
 /// these tests running at once would each reap the other's. One at a time.
@@ -108,3 +109,73 @@ fn pipe() -> (i32, i32) {
     assert_eq!(unsafe { libc::pipe(fds.as_mut_ptr()) }, 0);
     (fds[0], fds[1])
 }
+
+#[test]
+fn an_exit_reaches_whoever_asked_for_it() {
+    let _alone = alone();
+    let waiters = Waiters::new();
+
+    let mut exited = waiters
+        .watch(|| Ok(fork_child(|| unsafe { libc::_exit(9) })))
+        .expect("the child never started");
+
+    let deadline = Instant::now() + Duration::from_secs(5);
+    while Instant::now() < deadline {
+        for (pid, exit) in reap_exited() {
+            waiters.deliver(pid, exit);
+        }
+        if let Ok(exit) = exited.try_recv() {
+            assert_eq!(exit.exit_code, Some(9));
+            return;
+        }
+        std::thread::sleep(Duration::from_millis(10));
+    }
+    panic!("the exit was reaped by the reaper and never handed on");
+}
+
+#[test]
+fn an_exit_that_happens_before_the_caller_is_registered_is_not_lost() {
+    let _alone = alone();
+    let waiters = Waiters::new();
+
+    // A reaper already running, as it is in the guest: it will collect this
+    // child before `watch` has finished registering interest in it.
+    let stop = Arc::new(AtomicBool::new(false));
+    let reaper = std::thread::spawn({
+        let waiters = waiters.clone();
+        let stop = stop.clone();
+        move || {
+            while !stop.load(Ordering::Relaxed) {
+                for (pid, exit) in reap_exited() {
+                    waiters.deliver(pid, exit);
+                }
+                std::thread::sleep(Duration::from_millis(1));
+            }
+        }
+    });
+
+    let mut exited = waiters
+        .watch(|| {
+            let pid = fork_child(|| unsafe { libc::_exit(5) });
+            // Long enough for the child to exit and the reaper to reach it.
+            std::thread::sleep(Duration::from_millis(200));
+            Ok(pid)
+        })
+        .expect("the child never started");
+
+    let deadline = Instant::now() + Duration::from_secs(5);
+    let exit = loop {
+        if let Ok(exit) = exited.try_recv() {
+            break exit;
+        }
+        assert!(
+            Instant::now() < deadline,
+            "the exit was dropped on the way through"
+        );
+        std::thread::sleep(Duration::from_millis(10));
+    };
+    assert_eq!(exit.exit_code, Some(5));
+
+    stop.store(true, Ordering::Relaxed);
+    reaper.join().unwrap();
+}