mirror of
https://github.com/tokio-rs/tokio.git
synced 2026-09-02 00:00:11 +02:00
process: move into the tokio-net crate (#1475)
This commit is contained in:
Ivan Petkov
@@ -0,0 +1,13 @@
|
||||
use std::io;
|
||||
|
||||
/// An interface for killing a running process.
|
||||
pub(crate) trait Kill {
|
||||
/// Forcefully kill the process.
|
||||
fn kill(&mut self) -> io::Result<()>;
|
||||
}
|
||||
|
||||
impl<T: Kill> Kill for &mut T {
|
||||
fn kill(&mut self) -> io::Result<()> {
|
||||
(**self).kill()
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,225 @@
|
||||
//! Unix handling of child processes
|
||||
//!
|
||||
//! Right now the only "fancy" thing about this is how we implement the
|
||||
//! `Future` implementation on `Child` to get the exit status. Unix offers
|
||||
//! no way to register a child with epoll, and the only real way to get a
|
||||
//! notification when a process exits is the SIGCHLD signal.
|
||||
//!
|
||||
//! Signal handling in general is *super* hairy and complicated, and it's even
|
||||
//! more complicated here with the fact that signals are coalesced, so we may
|
||||
//! not get a SIGCHLD-per-child.
|
||||
//!
|
||||
//! Our best approximation here is to check *all spawned processes* for all
|
||||
//! SIGCHLD signals received. To do that we create a `Signal`, implemented in
|
||||
//! the `tokio-net` crate, which is a stream over signals being received.
|
||||
//!
|
||||
//! Later when we poll the process's exit status we simply check to see if a
|
||||
//! SIGCHLD has happened since we last checked, and while that returns "yes" we
|
||||
//! keep trying.
|
||||
//!
|
||||
//! Note that this means that this isn't really scalable, but then again
|
||||
//! processes in general aren't scalable (e.g. millions) so it shouldn't be that
|
||||
//! bad in theory...
|
||||
|
||||
mod orphan;
|
||||
mod reap;
|
||||
|
||||
use self::orphan::{AtomicOrphanQueue, OrphanQueue, Wait};
|
||||
use self::reap::Reaper;
|
||||
use super::SpawnedChild;
|
||||
use crate::process::kill::Kill;
|
||||
use crate::signal::unix::{signal, Signal, SignalKind};
|
||||
use crate::util::PollEvented;
|
||||
use mio::event::Evented;
|
||||
use mio::unix::{EventedFd, UnixReady};
|
||||
use mio::{Poll as MioPoll, PollOpt, Ready, Token};
|
||||
use std::fmt;
|
||||
use std::future::Future;
|
||||
use std::io;
|
||||
use std::os::unix::io::{AsRawFd, RawFd};
|
||||
use std::pin::Pin;
|
||||
use std::process::{self, ExitStatus};
|
||||
use std::task::Context;
|
||||
use std::task::Poll;
|
||||
|
||||
impl Wait for process::Child {
|
||||
fn id(&self) -> u32 {
|
||||
self.id()
|
||||
}
|
||||
|
||||
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()
|
||||
}
|
||||
}
|
||||
|
||||
lazy_static::lazy_static! {
|
||||
static ref ORPHAN_QUEUE: AtomicOrphanQueue<process::Child> = AtomicOrphanQueue::new();
|
||||
}
|
||||
|
||||
struct GlobalOrphanQueue;
|
||||
|
||||
impl fmt::Debug for GlobalOrphanQueue {
|
||||
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
ORPHAN_QUEUE.fmt(fmt)
|
||||
}
|
||||
}
|
||||
|
||||
impl OrphanQueue<process::Child> for GlobalOrphanQueue {
|
||||
fn push_orphan(&self, orphan: process::Child) {
|
||||
ORPHAN_QUEUE.push_orphan(orphan)
|
||||
}
|
||||
|
||||
fn reap_orphans(&self) {
|
||||
ORPHAN_QUEUE.reap_orphans()
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use = "futures do nothing unless polled"]
|
||||
pub(crate) struct Child {
|
||||
inner: Reaper<process::Child, GlobalOrphanQueue, Signal>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for Child {
|
||||
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
fmt.debug_struct("Child")
|
||||
.field("pid", &self.inner.id())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_child(cmd: &mut process::Command) -> io::Result<SpawnedChild> {
|
||||
let mut child = cmd.spawn()?;
|
||||
let stdin = stdio(child.stdin.take())?;
|
||||
let stdout = stdio(child.stdout.take())?;
|
||||
let stderr = stdio(child.stderr.take())?;
|
||||
|
||||
let signal = signal(SignalKind::child())?;
|
||||
|
||||
Ok(SpawnedChild {
|
||||
child: Child {
|
||||
inner: Reaper::new(child, GlobalOrphanQueue, signal),
|
||||
},
|
||||
stdin,
|
||||
stdout,
|
||||
stderr,
|
||||
})
|
||||
}
|
||||
|
||||
impl Child {
|
||||
pub(crate) fn id(&self) -> u32 {
|
||||
self.inner.id()
|
||||
}
|
||||
}
|
||||
|
||||
impl Kill for Child {
|
||||
fn kill(&mut self) -> io::Result<()> {
|
||||
self.inner.kill()
|
||||
}
|
||||
}
|
||||
|
||||
impl Future for Child {
|
||||
type Output = io::Result<ExitStatus>;
|
||||
|
||||
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||
Pin::new(&mut self.inner).poll(cx)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct Fd<T> {
|
||||
inner: T,
|
||||
}
|
||||
|
||||
impl<T> io::Read for Fd<T>
|
||||
where
|
||||
T: io::Read,
|
||||
{
|
||||
fn read(&mut self, bytes: &mut [u8]) -> io::Result<usize> {
|
||||
self.inner.read(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> io::Write for Fd<T>
|
||||
where
|
||||
T: io::Write,
|
||||
{
|
||||
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
|
||||
self.inner.write(bytes)
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.inner.flush()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> AsRawFd for Fd<T>
|
||||
where
|
||||
T: AsRawFd,
|
||||
{
|
||||
fn as_raw_fd(&self) -> RawFd {
|
||||
self.inner.as_raw_fd()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Evented for Fd<T>
|
||||
where
|
||||
T: AsRawFd,
|
||||
{
|
||||
fn register(
|
||||
&self,
|
||||
poll: &MioPoll,
|
||||
token: Token,
|
||||
interest: Ready,
|
||||
opts: PollOpt,
|
||||
) -> io::Result<()> {
|
||||
EventedFd(&self.as_raw_fd()).register(poll, token, interest | UnixReady::hup(), opts)
|
||||
}
|
||||
|
||||
fn reregister(
|
||||
&self,
|
||||
poll: &MioPoll,
|
||||
token: Token,
|
||||
interest: Ready,
|
||||
opts: PollOpt,
|
||||
) -> io::Result<()> {
|
||||
EventedFd(&self.as_raw_fd()).reregister(poll, token, interest | UnixReady::hup(), opts)
|
||||
}
|
||||
|
||||
fn deregister(&self, poll: &MioPoll) -> io::Result<()> {
|
||||
EventedFd(&self.as_raw_fd()).deregister(poll)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) type ChildStdin = PollEvented<Fd<process::ChildStdin>>;
|
||||
pub(crate) type ChildStdout = PollEvented<Fd<process::ChildStdout>>;
|
||||
pub(crate) type ChildStderr = PollEvented<Fd<process::ChildStderr>>;
|
||||
|
||||
fn stdio<T>(option: Option<T>) -> io::Result<Option<PollEvented<Fd<T>>>>
|
||||
where
|
||||
T: AsRawFd,
|
||||
{
|
||||
let io = match option {
|
||||
Some(io) => io,
|
||||
None => return Ok(None),
|
||||
};
|
||||
|
||||
// Set the fd to nonblocking before we pass it to the event loop
|
||||
unsafe {
|
||||
let fd = io.as_raw_fd();
|
||||
let r = libc::fcntl(fd, libc::F_GETFL);
|
||||
if r == -1 {
|
||||
return Err(io::Error::last_os_error());
|
||||
}
|
||||
let r = libc::fcntl(fd, libc::F_SETFL, r | libc::O_NONBLOCK);
|
||||
if r == -1 {
|
||||
return Err(io::Error::last_os_error());
|
||||
}
|
||||
}
|
||||
Ok(Some(PollEvented::new(Fd { inner: io })))
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
use crossbeam_queue::SegQueue;
|
||||
use log::error;
|
||||
use std::io;
|
||||
use std::process::ExitStatus;
|
||||
|
||||
/// An interface for waiting on a process to exit.
|
||||
pub(crate) trait Wait {
|
||||
/// Get the identifier for this process or diagnostics.
|
||||
fn id(&self) -> u32;
|
||||
/// Try waiting for a process to exit in a non-blocking manner.
|
||||
fn try_wait(&mut self) -> io::Result<Option<ExitStatus>>;
|
||||
}
|
||||
|
||||
impl<T: Wait> Wait for &mut T {
|
||||
fn id(&self) -> u32 {
|
||||
(**self).id()
|
||||
}
|
||||
|
||||
fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
|
||||
(**self).try_wait()
|
||||
}
|
||||
}
|
||||
|
||||
/// An interface for queueing up an orphaned process so that it can be reaped.
|
||||
pub(crate) trait OrphanQueue<T> {
|
||||
/// Add an orphan to the queue.
|
||||
fn push_orphan(&self, orphan: T);
|
||||
/// Attempt to reap every process in the queue, ignoring any errors and
|
||||
/// enqueueing any orphans which have not yet exited.
|
||||
fn reap_orphans(&self);
|
||||
}
|
||||
|
||||
impl<T, O: OrphanQueue<T>> OrphanQueue<T> for &O {
|
||||
fn push_orphan(&self, orphan: T) {
|
||||
(**self).push_orphan(orphan);
|
||||
}
|
||||
|
||||
fn reap_orphans(&self) {
|
||||
(**self).reap_orphans()
|
||||
}
|
||||
}
|
||||
|
||||
/// An atomic implementation of `OrphanQueue`.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct AtomicOrphanQueue<T> {
|
||||
queue: SegQueue<T>,
|
||||
}
|
||||
|
||||
impl<T> AtomicOrphanQueue<T> {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
queue: SegQueue::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Wait> OrphanQueue<T> for AtomicOrphanQueue<T> {
|
||||
fn push_orphan(&self, orphan: T) {
|
||||
self.queue.push(orphan)
|
||||
}
|
||||
|
||||
fn reap_orphans(&self) {
|
||||
let len = self.queue.len();
|
||||
|
||||
if len == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
let mut orphans = Vec::with_capacity(len);
|
||||
while let Ok(mut orphan) = self.queue.pop() {
|
||||
match orphan.try_wait() {
|
||||
Ok(Some(_)) => {}
|
||||
Err(e) => error!(
|
||||
"leaking orphaned process {} due to try_wait() error: {}",
|
||||
orphan.id(),
|
||||
e,
|
||||
),
|
||||
|
||||
// Still not done yet, we need to put it back in the queue
|
||||
// when were done draining it, so that we don't get stuck
|
||||
// in an infinite loop here
|
||||
Ok(None) => orphans.push(orphan),
|
||||
}
|
||||
}
|
||||
|
||||
for orphan in orphans {
|
||||
self.queue.push(orphan);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::Wait;
|
||||
use super::{AtomicOrphanQueue, OrphanQueue};
|
||||
use std::cell::Cell;
|
||||
use std::io;
|
||||
use std::os::unix::process::ExitStatusExt;
|
||||
use std::process::ExitStatus;
|
||||
use std::rc::Rc;
|
||||
|
||||
struct MockWait {
|
||||
total_waits: Rc<Cell<usize>>,
|
||||
num_wait_until_status: usize,
|
||||
return_err: bool,
|
||||
}
|
||||
|
||||
impl MockWait {
|
||||
fn new(num_wait_until_status: usize) -> Self {
|
||||
Self {
|
||||
total_waits: Rc::new(Cell::new(0)),
|
||||
num_wait_until_status,
|
||||
return_err: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn with_err() -> Self {
|
||||
Self {
|
||||
total_waits: Rc::new(Cell::new(0)),
|
||||
num_wait_until_status: 0,
|
||||
return_err: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Wait for MockWait {
|
||||
fn id(&self) -> u32 {
|
||||
42
|
||||
}
|
||||
|
||||
fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
|
||||
let waits = self.total_waits.get();
|
||||
|
||||
let ret = if self.num_wait_until_status == waits {
|
||||
if self.return_err {
|
||||
Ok(Some(ExitStatus::from_raw(0)))
|
||||
} else {
|
||||
Err(io::Error::new(io::ErrorKind::Other, "mock err"))
|
||||
}
|
||||
} else {
|
||||
Ok(None)
|
||||
};
|
||||
|
||||
self.total_waits.set(waits + 1);
|
||||
ret
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drain_attempts_a_single_reap_of_all_queued_orphans() {
|
||||
let first_orphan = MockWait::new(0);
|
||||
let second_orphan = MockWait::new(1);
|
||||
let third_orphan = MockWait::new(2);
|
||||
let fourth_orphan = MockWait::with_err();
|
||||
|
||||
let first_waits = first_orphan.total_waits.clone();
|
||||
let second_waits = second_orphan.total_waits.clone();
|
||||
let third_waits = third_orphan.total_waits.clone();
|
||||
let fourth_waits = fourth_orphan.total_waits.clone();
|
||||
|
||||
let orphanage = AtomicOrphanQueue::new();
|
||||
orphanage.push_orphan(first_orphan);
|
||||
orphanage.push_orphan(third_orphan);
|
||||
orphanage.push_orphan(second_orphan);
|
||||
orphanage.push_orphan(fourth_orphan);
|
||||
|
||||
assert_eq!(orphanage.queue.len(), 4);
|
||||
|
||||
orphanage.reap_orphans();
|
||||
assert_eq!(orphanage.queue.len(), 2);
|
||||
assert_eq!(first_waits.get(), 1);
|
||||
assert_eq!(second_waits.get(), 1);
|
||||
assert_eq!(third_waits.get(), 1);
|
||||
assert_eq!(fourth_waits.get(), 1);
|
||||
|
||||
orphanage.reap_orphans();
|
||||
assert_eq!(orphanage.queue.len(), 1);
|
||||
assert_eq!(first_waits.get(), 1);
|
||||
assert_eq!(second_waits.get(), 2);
|
||||
assert_eq!(third_waits.get(), 2);
|
||||
assert_eq!(fourth_waits.get(), 1);
|
||||
|
||||
orphanage.reap_orphans();
|
||||
assert_eq!(orphanage.queue.len(), 0);
|
||||
assert_eq!(first_waits.get(), 1);
|
||||
assert_eq!(second_waits.get(), 2);
|
||||
assert_eq!(third_waits.get(), 3);
|
||||
assert_eq!(fourth_waits.get(), 1);
|
||||
|
||||
orphanage.reap_orphans(); // Safe to reap when empty
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,334 @@
|
||||
use super::orphan::{OrphanQueue, Wait};
|
||||
use crate::process::kill::Kill;
|
||||
use futures_core::stream::Stream;
|
||||
use std::future::Future;
|
||||
use std::io;
|
||||
use std::ops::Deref;
|
||||
use std::pin::Pin;
|
||||
use std::process::ExitStatus;
|
||||
use std::task::Context;
|
||||
use std::task::Poll;
|
||||
|
||||
/// Orchestrates between registering interest for receiving signals when a
|
||||
/// child process has exited, and attempting to poll for process completion.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct Reaper<W, Q, S>
|
||||
where
|
||||
W: Wait + Unpin,
|
||||
Q: OrphanQueue<W>,
|
||||
{
|
||||
inner: Option<W>,
|
||||
orphan_queue: Q,
|
||||
signal: S,
|
||||
}
|
||||
|
||||
impl<W, Q, S> Deref for Reaper<W, Q, S>
|
||||
where
|
||||
W: Wait + Unpin,
|
||||
Q: OrphanQueue<W>,
|
||||
{
|
||||
type Target = W;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
self.inner()
|
||||
}
|
||||
}
|
||||
|
||||
impl<W, Q, S> Reaper<W, Q, S>
|
||||
where
|
||||
W: Wait + Unpin,
|
||||
Q: OrphanQueue<W>,
|
||||
{
|
||||
pub(crate) fn new(inner: W, orphan_queue: Q, signal: S) -> Self {
|
||||
Self {
|
||||
inner: Some(inner),
|
||||
orphan_queue,
|
||||
signal,
|
||||
}
|
||||
}
|
||||
|
||||
fn inner(&self) -> &W {
|
||||
self.inner.as_ref().expect("inner has gone away")
|
||||
}
|
||||
|
||||
fn inner_mut(&mut self) -> &mut W {
|
||||
self.inner.as_mut().expect("inner has gone away")
|
||||
}
|
||||
}
|
||||
|
||||
impl<W, Q, S> Future for Reaper<W, Q, S>
|
||||
where
|
||||
W: Wait + Unpin,
|
||||
Q: OrphanQueue<W> + Unpin,
|
||||
S: Stream + Unpin,
|
||||
{
|
||||
type Output = io::Result<ExitStatus>;
|
||||
|
||||
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||
loop {
|
||||
// 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 = Pin::new(&mut self.signal).poll_next(cx).is_pending();
|
||||
|
||||
self.orphan_queue.reap_orphans();
|
||||
if let Some(status) = self.inner_mut().try_wait()? {
|
||||
return Poll::Ready(Ok(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 Poll::Pending;
|
||||
} 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, Q, S> Kill for Reaper<W, Q, S>
|
||||
where
|
||||
W: Kill + Wait + Unpin,
|
||||
Q: OrphanQueue<W>,
|
||||
{
|
||||
fn kill(&mut self) -> io::Result<()> {
|
||||
self.inner_mut().kill()
|
||||
}
|
||||
}
|
||||
|
||||
impl<W, Q, S> Drop for Reaper<W, Q, S>
|
||||
where
|
||||
W: Wait + Unpin,
|
||||
Q: OrphanQueue<W>,
|
||||
{
|
||||
fn drop(&mut self) {
|
||||
if let Ok(Some(_)) = self.inner_mut().try_wait() {
|
||||
return;
|
||||
}
|
||||
|
||||
let orphan = self.inner.take().unwrap();
|
||||
self.orphan_queue.push_orphan(orphan);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
use futures_core::stream::Stream;
|
||||
use futures_util::future::FutureExt;
|
||||
use std::cell::{Cell, RefCell};
|
||||
use std::os::unix::process::ExitStatusExt;
|
||||
use std::pin::Pin;
|
||||
use std::process::ExitStatus;
|
||||
use std::task::Context;
|
||||
use std::task::Poll;
|
||||
|
||||
#[derive(Debug)]
|
||||
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 id(&self) -> u32 {
|
||||
0
|
||||
}
|
||||
|
||||
fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
|
||||
let ret = if self.num_wait_until_status == self.total_waits {
|
||||
Some(self.status)
|
||||
} 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 = io::Result<()>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
let inner = Pin::get_mut(self);
|
||||
inner.total_polls += 1;
|
||||
match inner.values.remove(0) {
|
||||
Some(()) => Poll::Ready(Some(Ok(()))),
|
||||
None => Poll::Pending,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct MockQueue<W> {
|
||||
all_enqueued: RefCell<Vec<W>>,
|
||||
total_reaps: Cell<usize>,
|
||||
}
|
||||
|
||||
impl<W> MockQueue<W> {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
all_enqueued: RefCell::new(Vec::new()),
|
||||
total_reaps: Cell::new(0),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<W: Wait> OrphanQueue<W> for MockQueue<W> {
|
||||
fn push_orphan(&self, orphan: W) {
|
||||
self.all_enqueued.borrow_mut().push(orphan);
|
||||
}
|
||||
|
||||
fn reap_orphans(&self) {
|
||||
self.total_reaps.set(self.total_reaps.get() + 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reaper() {
|
||||
let exit = ExitStatus::from_raw(0);
|
||||
let mock = MockWait::new(exit, 3);
|
||||
let mut grim = Reaper::new(
|
||||
mock,
|
||||
MockQueue::new(),
|
||||
MockStream::new(vec![None, Some(()), None, None, None]),
|
||||
);
|
||||
|
||||
let waker = futures_util::task::noop_waker();
|
||||
let mut context = Context::from_waker(&waker);
|
||||
|
||||
// Not yet exited, interest registered
|
||||
assert!(grim.poll_unpin(&mut context).is_pending());
|
||||
assert_eq!(1, grim.signal.total_polls);
|
||||
assert_eq!(1, grim.total_waits);
|
||||
assert_eq!(1, grim.orphan_queue.total_reaps.get());
|
||||
assert!(grim.orphan_queue.all_enqueued.borrow().is_empty());
|
||||
|
||||
// Not yet exited, couldn't register interest the first time
|
||||
// but managed to register interest the second time around
|
||||
assert!(grim.poll_unpin(&mut context).is_pending());
|
||||
assert_eq!(3, grim.signal.total_polls);
|
||||
assert_eq!(3, grim.total_waits);
|
||||
assert_eq!(3, grim.orphan_queue.total_reaps.get());
|
||||
assert!(grim.orphan_queue.all_enqueued.borrow().is_empty());
|
||||
|
||||
// Exited
|
||||
if let Poll::Ready(r) = grim.poll_unpin(&mut context) {
|
||||
assert!(r.is_ok());
|
||||
let exit_code = r.unwrap();
|
||||
assert_eq!(exit_code, exit);
|
||||
} else {
|
||||
unreachable!();
|
||||
}
|
||||
assert_eq!(4, grim.signal.total_polls);
|
||||
assert_eq!(4, grim.total_waits);
|
||||
assert_eq!(4, grim.orphan_queue.total_reaps.get());
|
||||
assert!(grim.orphan_queue.all_enqueued.borrow().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kill() {
|
||||
let exit = ExitStatus::from_raw(0);
|
||||
let mut grim = Reaper::new(
|
||||
MockWait::new(exit, 0),
|
||||
MockQueue::new(),
|
||||
MockStream::new(vec![None]),
|
||||
);
|
||||
|
||||
grim.kill().unwrap();
|
||||
assert_eq!(1, grim.total_kills);
|
||||
assert_eq!(0, grim.orphan_queue.total_reaps.get());
|
||||
assert!(grim.orphan_queue.all_enqueued.borrow().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_reaps_if_possible() {
|
||||
let exit = ExitStatus::from_raw(0);
|
||||
let mut mock = MockWait::new(exit, 0);
|
||||
|
||||
{
|
||||
let queue = MockQueue::new();
|
||||
|
||||
let grim = Reaper::new(&mut mock, &queue, MockStream::new(vec![]));
|
||||
|
||||
drop(grim);
|
||||
|
||||
assert_eq!(0, queue.total_reaps.get());
|
||||
assert!(queue.all_enqueued.borrow().is_empty());
|
||||
}
|
||||
|
||||
assert_eq!(1, mock.total_waits);
|
||||
assert_eq!(0, mock.total_kills);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_enqueues_orphan_if_wait_fails() {
|
||||
let exit = ExitStatus::from_raw(0);
|
||||
let mut mock = MockWait::new(exit, 2);
|
||||
|
||||
{
|
||||
let queue = MockQueue::<&mut MockWait>::new();
|
||||
let grim = Reaper::new(&mut mock, &queue, MockStream::new(vec![]));
|
||||
drop(grim);
|
||||
|
||||
assert_eq!(0, queue.total_reaps.get());
|
||||
assert_eq!(1, queue.all_enqueued.borrow().len());
|
||||
}
|
||||
|
||||
assert_eq!(1, mock.total_waits);
|
||||
assert_eq!(0, mock.total_kills);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
//! Windows asynchronous process handling.
|
||||
//!
|
||||
//! Like with Unix we don't actually have a way of registering a process with an
|
||||
//! IOCP object. As a result we similarly need another mechanism for getting a
|
||||
//! signal when a process has exited. For now this is implemented with the
|
||||
//! `RegisterWaitForSingleObject` function in the kernel32.dll.
|
||||
//!
|
||||
//! This strategy is the same that libuv takes and essentially just queues up a
|
||||
//! wait for the process in a kernel32-specific thread pool. Once the object is
|
||||
//! notified (e.g. the process exits) then we have a callback that basically
|
||||
//! just completes a `Oneshot`.
|
||||
//!
|
||||
//! The `poll_exit` implementation will attempt to wait for the process in a
|
||||
//! nonblocking fashion, but failing that it'll fire off a
|
||||
//! `RegisterWaitForSingleObject` and then wait on the other end of the oneshot
|
||||
//! from then on out.
|
||||
|
||||
use super::SpawnedChild;
|
||||
use crate::process::kill::Kill;
|
||||
use crate::util::PollEvented;
|
||||
use futures_util::future::Fuse;
|
||||
use futures_util::future::FutureExt;
|
||||
use mio_named_pipes::NamedPipe;
|
||||
use std::fmt;
|
||||
use std::future::Future;
|
||||
use std::io;
|
||||
use std::os::windows::prelude::*;
|
||||
use std::os::windows::process::ExitStatusExt;
|
||||
use std::pin::Pin;
|
||||
use std::process::{Child as StdChild, Command as StdCommand, ExitStatus};
|
||||
use std::ptr;
|
||||
use std::task::Context;
|
||||
use std::task::Poll;
|
||||
use tokio_sync::oneshot;
|
||||
use winapi::shared::minwindef::FALSE;
|
||||
use winapi::shared::winerror::WAIT_TIMEOUT;
|
||||
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
|
||||
use winapi::um::processthreadsapi::GetExitCodeProcess;
|
||||
use winapi::um::synchapi::WaitForSingleObject;
|
||||
use winapi::um::threadpoollegacyapiset::UnregisterWaitEx;
|
||||
use winapi::um::winbase::{RegisterWaitForSingleObject, INFINITE, WAIT_OBJECT_0};
|
||||
use winapi::um::winnt::{BOOLEAN, HANDLE, PVOID, WT_EXECUTEINWAITTHREAD, WT_EXECUTEONLYONCE};
|
||||
|
||||
#[must_use = "futures do nothing unless polled"]
|
||||
pub(crate) struct Child {
|
||||
child: StdChild,
|
||||
waiting: Option<Waiting>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for Child {
|
||||
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
fmt.debug_struct("Child")
|
||||
.field("pid", &self.id())
|
||||
.field("child", &self.child)
|
||||
.field("waiting", &"..")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
struct Waiting {
|
||||
rx: Fuse<oneshot::Receiver<()>>,
|
||||
wait_object: HANDLE,
|
||||
tx: *mut Option<oneshot::Sender<()>>,
|
||||
}
|
||||
|
||||
unsafe impl Sync for Waiting {}
|
||||
unsafe impl Send for Waiting {}
|
||||
|
||||
pub(crate) fn spawn_child(cmd: &mut StdCommand) -> io::Result<SpawnedChild> {
|
||||
let mut child = cmd.spawn()?;
|
||||
let stdin = stdio(child.stdin.take());
|
||||
let stdout = stdio(child.stdout.take());
|
||||
let stderr = stdio(child.stderr.take());
|
||||
|
||||
Ok(SpawnedChild {
|
||||
child: Child {
|
||||
child,
|
||||
waiting: None,
|
||||
},
|
||||
stdin,
|
||||
stdout,
|
||||
stderr,
|
||||
})
|
||||
}
|
||||
|
||||
impl Child {
|
||||
pub(crate) fn id(&self) -> u32 {
|
||||
self.child.id()
|
||||
}
|
||||
}
|
||||
|
||||
impl Kill for Child {
|
||||
fn kill(&mut self) -> io::Result<()> {
|
||||
self.child.kill()
|
||||
}
|
||||
}
|
||||
|
||||
impl Future for Child {
|
||||
type Output = io::Result<ExitStatus>;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||
let inner = Pin::get_mut(self);
|
||||
loop {
|
||||
if let Some(ref mut w) = inner.waiting {
|
||||
match w.rx.poll_unpin(cx) {
|
||||
Poll::Ready(Ok(())) => {}
|
||||
Poll::Ready(Err(_)) => panic!("should not be canceled"),
|
||||
Poll::Pending => return Poll::Pending,
|
||||
}
|
||||
let status = try_wait(&inner.child)?.expect("not ready yet");
|
||||
return Poll::Ready(Ok(status.into()));
|
||||
}
|
||||
|
||||
if let Some(e) = try_wait(&inner.child)? {
|
||||
return Poll::Ready(Ok(e.into()));
|
||||
}
|
||||
let (tx, rx) = oneshot::channel();
|
||||
let ptr = Box::into_raw(Box::new(Some(tx)));
|
||||
let mut wait_object = ptr::null_mut();
|
||||
let rc = unsafe {
|
||||
RegisterWaitForSingleObject(
|
||||
&mut wait_object,
|
||||
inner.child.as_raw_handle(),
|
||||
Some(callback),
|
||||
ptr as *mut _,
|
||||
INFINITE,
|
||||
WT_EXECUTEINWAITTHREAD | WT_EXECUTEONLYONCE,
|
||||
)
|
||||
};
|
||||
if rc == 0 {
|
||||
let err = io::Error::last_os_error();
|
||||
drop(unsafe { Box::from_raw(ptr) });
|
||||
return Poll::Ready(Err(err));
|
||||
}
|
||||
inner.waiting = Some(Waiting {
|
||||
rx: rx.fuse(),
|
||||
wait_object,
|
||||
tx: ptr,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Waiting {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
let rc = UnregisterWaitEx(self.wait_object, INVALID_HANDLE_VALUE);
|
||||
if rc == 0 {
|
||||
panic!("failed to unregister: {}", io::Error::last_os_error());
|
||||
}
|
||||
drop(Box::from_raw(self.tx));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe extern "system" fn callback(ptr: PVOID, _timer_fired: BOOLEAN) {
|
||||
let complete = &mut *(ptr as *mut Option<oneshot::Sender<()>>);
|
||||
let _ = complete.take().unwrap().send(());
|
||||
}
|
||||
|
||||
pub(crate) fn try_wait(child: &StdChild) -> io::Result<Option<ExitStatus>> {
|
||||
unsafe {
|
||||
match WaitForSingleObject(child.as_raw_handle(), 0) {
|
||||
WAIT_OBJECT_0 => {}
|
||||
WAIT_TIMEOUT => return Ok(None),
|
||||
_ => return Err(io::Error::last_os_error()),
|
||||
}
|
||||
let mut status = 0;
|
||||
let rc = GetExitCodeProcess(child.as_raw_handle(), &mut status);
|
||||
if rc == FALSE {
|
||||
Err(io::Error::last_os_error())
|
||||
} else {
|
||||
Ok(Some(ExitStatus::from_raw(status)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) type ChildStdin = PollEvented<NamedPipe>;
|
||||
pub(crate) type ChildStdout = PollEvented<NamedPipe>;
|
||||
pub(crate) type ChildStderr = PollEvented<NamedPipe>;
|
||||
|
||||
fn stdio<T>(option: Option<T>) -> Option<PollEvented<NamedPipe>>
|
||||
where
|
||||
T: IntoRawHandle,
|
||||
{
|
||||
let io = match option {
|
||||
Some(io) => io,
|
||||
None => return None,
|
||||
};
|
||||
let pipe = unsafe { NamedPipe::from_raw_handle(io.into_raw_handle()) };
|
||||
Some(PollEvented::new(pipe))
|
||||
}
|
||||
Reference in New Issue
Block a user