process: Refactor Unix process handling

src/unix/mod.rs

@@ -24,35 +24,44 @@
 extern crate libc;
 extern crate tokio_signal;
 
-use std::io;
+mod reap;
-use std::os::unix::prelude::*;
-use std::process::{self, ExitStatus};
 
 use futures::future::FlattenStream;
-use futures::{Future, Poll, Async, Stream};
+use futures::{Future, Poll};
 use mio::unix::{EventedFd, UnixReady};
 use mio::{PollOpt, Ready, Token};
 use mio::event::Evented;
 use mio;
+use self::reap::{EventedReaper, Kill, Wait};
 use self::tokio_signal::unix::Signal;
 use std::fmt;
+use std::io;
+use std::os::unix::io::{AsRawFd, RawFd};
+use std::process::{self, ExitStatus};
 use tokio_io::IoFuture;
 use tokio_reactor::{Handle, PollEvented};
 
+impl Wait for process::Child {
+    fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
+        self.try_wait()
+    }
+}
+
+impl Kill for process::Child {
+    fn kill(&mut self) -> io::Result<()> {
+        self.kill()
+    }
+}
+
 #[must_use = "futures do nothing unless polled"]
 pub struct Child {
-    inner: process::Child,
+    inner: EventedReaper<process::Child, FlattenStream<IoFuture<Signal>>>,
-    reaped: bool,
-    sigchld: FlattenStream<IoFuture<Signal>>,
 }
 
 impl fmt::Debug for Child {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         fmt.debug_struct("Child")
             .field("pid", &self.inner.id())
-            .field("inner", &self.inner)
-            .field("reaped", &self.reaped)
-            .field("sigchld", &"..")
             .finish()
     }
 }
@@ -65,10 +74,9 @@ impl Child {
         let stdout = stdio(inner.stdout.take(), handle)?;
         let stderr = stdio(inner.stderr.take(), handle)?;
 
+        let signal = Signal::with_handle(libc::SIGCHLD, handle).flatten_stream();
         let child = Child {
-            inner: inner,
+            inner: EventedReaper::new(inner, signal),
-            reaped: false,
-            sigchld: Signal::with_handle(libc::SIGCHLD, handle).flatten_stream(),
         };
 
         Ok((child, stdin, stdout, stderr))
@@ -79,93 +87,11 @@ impl Child {
     }
 
     pub fn kill(&mut self) -> io::Result<()> {
-        if !self.reaped {
+        self.inner.kill()
-            // NB: SIGKILL cannnot be caught, so the process will definitely exit immediately.
-            // We're not waiting for the process itself but for the kernel to execute the kill.
-            self.inner.kill()?;
-            let _ = self.try_wait(true);
-        }
-
-        Ok(())
     }
 
     pub fn poll_exit(&mut self) -> Poll<ExitStatus, io::Error> {
-        loop {
+        self.inner.poll()
-            // Ensure we don't register for additional notifications
-            // if the child has already finished.
-            if self.reaped {
-                return Ok(Async::NotReady);
-            }
-
-            // If the child hasn't exited yet, then it's our responsibility to
-            // ensure the current task gets notified when it might be able to
-            // make progress.
-            //
-            // As described in `spawn` above, we just indicate that we can
-            // next make progress once a SIGCHLD is received.
-            //
-            // However, we will register for a notification on the next signal
-            // BEFORE we poll the child. Otherwise it is possible that the child
-            // can exit and the signal can arrive after we last polled the child,
-            // but before we've registered for a notification on the next signal
-            // (this can cause a deadlock if there are no more spawned children
-            // which can generate a different signal for us). A side effect of
-            // pre-registering for signal notifications is that when the child
-            // exits, we will have already registered for an additional
-            // notification we don't need to consume. If another signal arrives,
-            // this future's task will be notified/woken up again. Since the
-            // futures model allows for spurious wake ups this extra wakeup
-            // should not cause significant issues with parent futures.
-            let registered_interest = try!(self.sigchld.poll()).is_not_ready();
-
-            if let Some(e) = try!(self.try_wait(false)) {
-                return Ok(e.into());
-            }
-
-            // If our attempt to poll for the next signal was not ready, then
-            // we've arranged for our task to get notified and we can bail out.
-            if registered_interest {
-                return Ok(Async::NotReady);
-            } else {
-                // Otherwise, if the signal stream delivered a signal to us, we
-                // won't get notified at the next signal, so we'll loop and try
-                // again.
-                continue;
-            }
-        }
-    }
-
-    fn try_wait(&mut self, block_on_wait: bool) -> io::Result<Option<ExitStatus>> {
-        assert!(!self.reaped);
-        let exit = try!(try_wait_process(self.id() as libc::pid_t, block_on_wait));
-
-        if let Some(_) = exit {
-            self.reaped = true;
-        }
-
-        Ok(exit)
-    }
-}
-
-fn try_wait_process(id: libc::pid_t, block_on_wait: bool) -> io::Result<Option<ExitStatus>> {
-    let wait_flags = if block_on_wait { 0 } else { libc::WNOHANG };
-    let mut status = 0;
-
-    loop {
-        match unsafe { libc::waitpid(id, &mut status, wait_flags) } {
-            0 => return Ok(None),
-            n if n < 0 => {
-                let err = io::Error::last_os_error();
-                if err.kind() == io::ErrorKind::Interrupted {
-                    continue
-                }
-                return Err(err)
-            }
-            n => {
-                assert_eq!(n, id);
-                return Ok(Some(ExitStatus::from_raw(status)))
-            }
-        }
     }
 }
 

src/unix/reap.rs

@@ -0,0 +1,313 @@
+use futures::{Async, Future, Poll, Stream};
+use std::io;
+use std::ops::Deref;
+use std::process::ExitStatus;
+
+/// An interface for waiting on a process to exit.
+pub trait Wait {
+    /// Try waiting for a process to exit in a non-blocking manner.
+    fn try_wait(&mut self) -> io::Result<Option<ExitStatus>>;
+}
+
+/// An interface for killing a running process.
+pub trait Kill {
+    /// Forcefully kill the process.
+    fn kill(&mut self) -> io::Result<()>;
+}
+
+#[derive(Debug, PartialEq)]
+enum WaitResult {
+    Exited(ExitStatus),
+    Reaped,
+}
+
+/// An interface for safely reaping a child process.
+trait Reap {
+    /// Try to reap the child process if ready.
+    fn try_reap(&mut self) -> Poll<WaitResult, io::Error>;
+}
+
+#[derive(Debug)]
+struct Reaper<W> {
+    reaped: bool,
+    proc: W,
+}
+
+impl<W> Reaper<W> {
+    fn new(proc: W) -> Self {
+        Self {
+            reaped: false,
+            proc,
+        }
+    }
+
+    fn reaped(&self) -> bool {
+        self.reaped
+    }
+}
+
+impl<W> Deref for Reaper<W> {
+    type Target = W;
+
+    fn deref(&self) -> &Self::Target {
+        &self.proc
+    }
+}
+
+impl<W: Wait> Reap for Reaper<W> {
+    fn try_reap(&mut self) -> Poll<WaitResult, io::Error> {
+        if self.reaped {
+            return Ok(Async::Ready(WaitResult::Reaped));
+        }
+
+        match self.proc.try_wait()? {
+            Some(exit) => {
+                self.reaped = true;
+                Ok(Async::Ready(WaitResult::Exited(exit)))
+            },
+            None => Ok(Async::NotReady),
+        }
+    }
+}
+
+impl<W: Kill> Kill for Reaper<W> {
+    fn kill(&mut self) -> io::Result<()> {
+        // NB: ensure we don't issue a kill after we've reaped the child
+        // since its process identifier could have been reused.
+        if self.reaped {
+            Ok(())
+        } else {
+            self.proc.kill()
+        }
+    }
+}
+
+/// Orchestrates between registering interest for receiving signals when a
+/// child process has exited, and attempting to poll for process completion.
+#[derive(Debug)]
+pub struct EventedReaper<W, S> {
+    inner: Reaper<W>,
+    signal: S,
+}
+
+impl<W, S> Deref for EventedReaper<W, S> {
+    type Target = W;
+
+    fn deref(&self) -> &Self::Target {
+        &*self.inner
+    }
+}
+
+impl<W, S> EventedReaper<W, S> {
+    pub fn new(inner: W, signal: S) -> Self {
+        Self {
+            inner: Reaper::new(inner),
+            signal,
+        }
+    }
+}
+
+impl<W, S> Future for EventedReaper<W, S>
+    where W: Wait,
+          S: Stream<Error = io::Error>,
+{
+    type Item = ExitStatus;
+    type Error = io::Error;
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        loop {
+            // Ensure we don't register for additional notifications
+            // if the child has already finished.
+            if self.inner.reaped() {
+                return Ok(Async::NotReady);
+            }
+
+            // If the child hasn't exited yet, then it's our responsibility to
+            // ensure the current task gets notified when it might be able to
+            // make progress.
+            //
+            // As described in `spawn` above, we just indicate that we can
+            // next make progress once a SIGCHLD is received.
+            //
+            // However, we will register for a notification on the next signal
+            // BEFORE we poll the child. Otherwise it is possible that the child
+            // can exit and the signal can arrive after we last polled the child,
+            // but before we've registered for a notification on the next signal
+            // (this can cause a deadlock if there are no more spawned children
+            // which can generate a different signal for us). A side effect of
+            // pre-registering for signal notifications is that when the child
+            // exits, we will have already registered for an additional
+            // notification we don't need to consume. If another signal arrives,
+            // this future's task will be notified/woken up again. Since the
+            // futures model allows for spurious wake ups this extra wakeup
+            // should not cause significant issues with parent futures.
+            let registered_interest = self.signal.poll()?.is_not_ready();
+
+            if let Async::Ready(WaitResult::Exited(status)) = self.inner.try_reap()? {
+                return Ok(Async::Ready(status));
+            }
+
+            // If our attempt to poll for the next signal was not ready, then
+            // we've arranged for our task to get notified and we can bail out.
+            if registered_interest {
+                return Ok(Async::NotReady);
+            } else {
+                // Otherwise, if the signal stream delivered a signal to us, we
+                // won't get notified at the next signal, so we'll loop and try
+                // again.
+                continue;
+            }
+        }
+    }
+}
+
+impl<W, S> Kill for EventedReaper<W, S>
+    where W: Kill,
+{
+    fn kill(&mut self) -> io::Result<()> {
+        self.inner.kill()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use futures::{Async, Poll, Stream};
+    use std::process::ExitStatus;
+    use std::os::unix::process::ExitStatusExt;
+    use super::*;
+
+    struct MockWait {
+        total_kills: usize,
+        total_waits: usize,
+        num_wait_until_status: usize,
+        status: ExitStatus,
+    }
+
+    impl MockWait {
+        fn new(status: ExitStatus, num_wait_until_status: usize) -> Self {
+            Self {
+                total_kills: 0,
+                total_waits: 0,
+                num_wait_until_status,
+                status
+            }
+        }
+    }
+
+    impl Wait for MockWait {
+        fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
+            let ret = if self.num_wait_until_status == self.total_waits {
+                Some(self.status.clone())
+            } else {
+                None
+            };
+
+            self.total_waits += 1;
+            Ok(ret)
+        }
+    }
+
+    impl Kill for MockWait {
+        fn kill(&mut self) -> io::Result<()> {
+            self.total_kills += 1;
+            Ok(())
+        }
+    }
+
+    struct MockStream {
+        total_polls: usize,
+        values: Vec<Option<()>>,
+    }
+
+    impl MockStream {
+        fn new(values: Vec<Option<()>>) -> Self {
+            Self {
+                total_polls: 0,
+                values
+            }
+        }
+    }
+
+    impl Stream for MockStream {
+        type Item = ();
+        type Error = io::Error;
+
+        fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+            self.total_polls += 1;
+            match self.values.remove(0) {
+                Some(()) => Ok(Async::Ready(Some(()))),
+                None => Ok(Async::NotReady),
+            }
+        }
+    }
+
+    #[test]
+    fn reaper() {
+        let exit = ExitStatus::from_raw(0);
+        let mock = MockWait::new(exit.clone(), 1);
+        let mut grim = Reaper::new(mock);
+
+        // Not yet exited
+        assert_eq!(Async::NotReady, grim.try_reap().expect("failed to wait"));
+        assert_eq!(1, grim.total_waits);
+
+        // Exited
+        assert_eq!(Async::Ready(WaitResult::Exited(exit)), grim.try_reap().expect("failed to wait"));
+        assert_eq!(2, grim.total_waits);
+
+        // Cannot call wait another time
+        assert_eq!(Async::Ready(WaitResult::Reaped), grim.try_reap().expect("failed to wait"));
+        assert_eq!(2, grim.total_waits);
+    }
+
+    #[test]
+    fn evented_reaper() {
+        let exit = ExitStatus::from_raw(0);
+        let mock = MockWait::new(exit.clone(), 3);
+        let mut grim = EventedReaper::new(mock, MockStream::new(vec!(
+            None,
+            Some(()),
+            None,
+            None,
+            None,
+        )));
+
+        // Not yet exited, interest registered
+        assert_eq!(Async::NotReady, grim.poll().expect("failed to wait"));
+        assert_eq!(1, grim.signal.total_polls);
+        assert_eq!(1, grim.total_waits);
+
+        // Not yet exited, couldn't register interest the first time
+        // but managed to register interest the second time around
+        assert_eq!(Async::NotReady, grim.poll().expect("failed to wait"));
+        assert_eq!(3, grim.signal.total_polls);
+        assert_eq!(3, grim.total_waits);
+
+        // Exited
+        assert_eq!(Async::Ready(exit), grim.poll().expect("failed to wait"));
+        assert_eq!(4, grim.signal.total_polls);
+        assert_eq!(4, grim.total_waits);
+
+        // Already reaped, no further calls
+        assert_eq!(Async::NotReady, grim.poll().expect("failed to poll"));
+        assert_eq!(4, grim.signal.total_polls);
+        assert_eq!(4, grim.total_waits);
+    }
+
+    #[test]
+    fn kill() {
+        let exit = ExitStatus::from_raw(0);
+        let mut grim = EventedReaper::new(
+            MockWait::new(exit, 0),
+            MockStream::new(vec!(None))
+        );
+
+        grim.kill().unwrap();
+        assert_eq!(1, grim.total_kills);
+
+        // Do not kill after reaping
+        assert_eq!(Async::Ready(exit), grim.poll().expect("failed to poll"));
+        grim.kill().unwrap();
+        assert_eq!(1, grim.total_kills);
+    }
+}