From 849771ecfa1e22fdd4f0bd299d10f0026ce14ed5 Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Tue, 12 Dec 2017 15:19:39 -0600 Subject: [PATCH] Add a `Handle::wakeup` method (#59) This method is intended to be used to wake up the reactor from a remote thread if necessary, forcing it to return from a blocked call of `turn` or otherwise prevent the next call to `turn` to from blocking. --- src/reactor/global.rs | 21 ++++++++++++++++----- src/reactor/mod.rs | 39 ++++++++++++++++++++++++++++++++++++--- tests/wakeup.rs | 43 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 95 insertions(+), 8 deletions(-) create mode 100644 tests/wakeup.rs diff --git a/src/reactor/global.rs b/src/reactor/global.rs index 69156d435..cd0f84e76 100644 --- a/src/reactor/global.rs +++ b/src/reactor/global.rs @@ -1,22 +1,28 @@ use std::io; use std::thread; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use reactor::{Reactor, Handle}; pub struct HelperThread { thread: Option>, reactor: Handle, + done: Arc, } impl HelperThread { pub fn new() -> io::Result { let reactor = Reactor::new()?; let reactor_handle = reactor.handle().clone(); - let thread = thread::Builder::new().spawn(move || run(reactor))?; + let done = Arc::new(AtomicBool::new(false)); + let done2 = done.clone(); + let thread = thread::Builder::new().spawn(move || run(reactor, done))?; Ok(HelperThread { thread: Some(thread), reactor: reactor_handle, + done: done2, }) } @@ -31,13 +37,18 @@ impl HelperThread { impl Drop for HelperThread { fn drop(&mut self) { - // TODO: kill the reactor thread and wait for it to exit, needs - // `Handle::wakeup` to be implemented in a future PR + let thread = match self.thread.take() { + Some(thread) => thread, + None => return + }; + self.done.store(true, Ordering::SeqCst); + self.reactor.wakeup(); + thread.join().unwrap(); } } -fn run(mut reactor: Reactor) { - loop { +fn run(mut reactor: Reactor, done: Arc) { + while !done.load(Ordering::SeqCst) { reactor.turn(None); } } diff --git a/src/reactor/mod.rs b/src/reactor/mod.rs index 5687b3a50..f794fd4f3 100644 --- a/src/reactor/mod.rs +++ b/src/reactor/mod.rs @@ -51,6 +51,8 @@ pub struct Reactor { /// State shared between the reactor and the handles. inner: Arc, + + _wakeup_registration: mio::Registration, } struct Inner { @@ -59,6 +61,9 @@ struct Inner { /// Dispatch slabs for I/O and futures events io_dispatch: RwLock>, + + /// Used to wake up the reactor from a call to `turn` + wakeup: mio::SetReadiness } /// A handle to an event loop. @@ -82,6 +87,7 @@ enum Direction { Write, } +const TOKEN_WAKEUP: mio::Token = mio::Token(0); const TOKEN_START: usize = 1; fn _assert_kinds() { @@ -94,14 +100,21 @@ impl Reactor { /// Creates a new event loop, returning any error that happened during the /// creation. pub fn new() -> io::Result { - // Create the I/O poller - let io = try!(mio::Poll::new()); + let io = mio::Poll::new()?; + let wakeup_pair = mio::Registration::new2(); + + io.register(&wakeup_pair.0, + TOKEN_WAKEUP, + mio::Ready::readable(), + mio::PollOpt::level())?; Ok(Reactor { events: mio::Events::with_capacity(1024), + _wakeup_registration: wakeup_pair.0, inner: Arc::new(Inner { io: io, io_dispatch: RwLock::new(Slab::with_capacity(1)), + wakeup: wakeup_pair.1, }), }) } @@ -146,7 +159,12 @@ impl Reactor { let token = event.token(); trace!("event {:?} {:?}", event.readiness(), event.token()); - self.dispatch(token, event.readiness()); + + if token == TOKEN_WAKEUP { + self.inner.wakeup.set_readiness(mio::Ready::empty()).unwrap(); + } else { + self.dispatch(token, event.readiness()); + } } } @@ -281,6 +299,21 @@ impl Handle { } } + /// Forces a reactor blocked in a call to `turn` to wakeup, or otherwise + /// makes the next call to `turn` return immediately. + /// + /// This method is intended to be used in situations where a notification + /// needs to otherwise be sent to the main reactor. If the reactor is + /// currently blocked inside of `turn` then it will wake up and soon return + /// after this method has been called. If the reactor is not currently + /// blocked in `turn`, then the next call to `turn` will not block and + /// return immediately. + pub fn wakeup(&self) { + if let Some(inner) = self.inner() { + inner.wakeup.set_readiness(mio::Ready::readable()).unwrap(); + } + } + fn into_usize(self) -> usize { unsafe { mem::transmute::, usize>(self.inner) diff --git a/tests/wakeup.rs b/tests/wakeup.rs new file mode 100644 index 000000000..490e6491f --- /dev/null +++ b/tests/wakeup.rs @@ -0,0 +1,43 @@ +extern crate tokio; + +use std::thread; +use std::time::{Duration, Instant}; +use std::sync::mpsc; + +use tokio::reactor::Reactor; + +#[test] +fn works() { + let mut r = Reactor::new().unwrap(); + r.handle().wakeup(); + r.turn(None); + + let now = Instant::now(); + let mut n = 0; + while now.elapsed() < Duration::from_millis(10) { + n += 1; + r.turn(Some(Duration::from_millis(10))); + } + assert!(n < 5); +} + +#[test] +fn wakes() { + const N: usize = 1_000; + + let mut r = Reactor::new().unwrap(); + let handle = r.handle(); + let (tx, rx) = mpsc::channel(); + let t = thread::spawn(move || { + for _ in 0..N { + rx.recv().unwrap(); + handle.wakeup(); + } + }); + + for _ in 0..N { + tx.send(()).unwrap(); + r.turn(None); + } + t.join().unwrap(); +}