New scheduler resilient to blocking

Author: Stjepan Glavina

Cargo.toml

@@ -26,6 +26,7 @@ default = [
   "async-task",
   "crossbeam-channel",
   "crossbeam-deque",
+  "crossbeam-queue",
   "futures-timer",
   "kv-log-macro",
   "log",
@@ -56,6 +57,7 @@ async-task = { version = "1.0.0", optional = true }
 broadcaster = { version = "0.2.6", optional = true, default-features = false, features = ["default-channels"] }
 crossbeam-channel = { version = "0.4.0", optional = true }
 crossbeam-deque = { version = "0.7.2", optional = true }
+crossbeam-queue = { version = "0.2.0", optional = true }
 crossbeam-utils = { version = "0.7.0", optional = true }
 futures-core = { version = "0.3.1", optional = true }
 futures-io = { version = "0.3.1", optional = true }

src/lib.rs

@@ -246,6 +246,8 @@ cfg_std! {
     pub mod stream;
     pub mod sync;
     pub mod task;
+
+    pub(crate) mod rt;
 }
 
 cfg_default! {

src/net/mod.rs

@@ -66,6 +66,5 @@ pub use tcp::{Incoming, TcpListener, TcpStream};
 pub use udp::UdpSocket;
 
 mod addr;
-pub(crate) mod driver;
 mod tcp;
 mod udp;

src/net/tcp/listener.rs

@@ -4,7 +4,7 @@ use std::pin::Pin;
 
 use crate::future;
 use crate::io;
-use crate::net::driver::Watcher;
+use crate::rt::Watcher;
 use crate::net::{TcpStream, ToSocketAddrs};
 use crate::stream::Stream;
 use crate::task::{Context, Poll};

src/net/tcp/stream.rs

@@ -4,7 +4,7 @@ use std::pin::Pin;
 
 use crate::future;
 use crate::io::{self, Read, Write};
-use crate::net::driver::Watcher;
+use crate::rt::Watcher;
 use crate::net::ToSocketAddrs;
 use crate::task::{spawn_blocking, Context, Poll};
 use crate::utils::Context as _;

src/net/udp/mod.rs

@@ -3,8 +3,8 @@ use std::net::SocketAddr;
 use std::net::{Ipv4Addr, Ipv6Addr};
 
 use crate::future;
-use crate::net::driver::Watcher;
 use crate::net::ToSocketAddrs;
+use crate::rt::Watcher;
 use crate::utils::Context as _;
 
 /// A UDP socket.
@@ -102,7 +102,7 @@ impl UdpSocket {
     /// ```no_run
     /// # fn main() -> std::io::Result<()> { async_std::task::block_on(async {
     /// #
-    ///	use async_std::net::UdpSocket;
+    /// use async_std::net::UdpSocket;
     ///
     /// let socket = UdpSocket::bind("127.0.0.1:0").await?;
     /// let addr = socket.local_addr()?;

src/os/unix/net/datagram.rs

@@ -8,7 +8,7 @@ use mio_uds;
 use super::SocketAddr;
 use crate::future;
 use crate::io;
-use crate::net::driver::Watcher;
+use crate::rt::Watcher;
 use crate::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
 use crate::path::Path;
 use crate::task::spawn_blocking;

src/os/unix/net/listener.rs

@@ -10,7 +10,7 @@ use super::SocketAddr;
 use super::UnixStream;
 use crate::future;
 use crate::io;
-use crate::net::driver::Watcher;
+use crate::rt::Watcher;
 use crate::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
 use crate::path::Path;
 use crate::stream::Stream;

src/os/unix/net/stream.rs

@@ -9,7 +9,7 @@ use mio_uds;
 
 use super::SocketAddr;
 use crate::io::{self, Read, Write};
-use crate::net::driver::Watcher;
+use crate::rt::Watcher;
 use crate::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
 use crate::path::Path;
 use crate::task::{spawn_blocking, Context, Poll};

src/rt/mod.rs

@@ -0,0 +1,23 @@
+//! The runtime.
+
+use std::thread;
+
+use once_cell::sync::Lazy;
+
+use crate::utils::abort_on_panic;
+
+pub use reactor::{Reactor, Watcher};
+pub use runtime::Runtime;
+
+mod reactor;
+mod runtime;
+
+/// The global runtime.
+pub static RUNTIME: Lazy<Runtime> = Lazy::new(|| {
+    thread::Builder::new()
+        .name("async-std/runtime".to_string())
+        .spawn(|| abort_on_panic(|| RUNTIME.run()))
+        .expect("cannot start a runtime thread");
+
+    Runtime::new()
+});

src/net/driver/mod.rs renamed to src/rt/reactor.rs

@@ -1,13 +1,13 @@
 use std::fmt;
 use std::sync::{Arc, Mutex};
+use std::time::Duration;
 
 use mio::{self, Evented};
-use once_cell::sync::Lazy;
 use slab::Slab;
 
 use crate::io;
+use crate::rt::RUNTIME;
 use crate::task::{Context, Poll, Waker};
-use crate::utils::abort_on_panic;
 
 /// Data associated with a registered I/O handle.
 #[derive(Debug)]
@@ -18,15 +18,18 @@ struct Entry {
     /// Tasks that are blocked on reading from this I/O handle.
     readers: Mutex<Vec<Waker>>,
 
-    /// Thasks that are blocked on writing to this I/O handle.
+    /// Tasks that are blocked on writing to this I/O handle.
     writers: Mutex<Vec<Waker>>,
 }
 
 /// The state of a networking driver.
-struct Reactor {
+pub struct Reactor {
     /// A mio instance that polls for new events.
     poller: mio::Poll,
 
+    /// A list into which mio stores events.
+    events: Mutex<mio::Events>,
+
     /// A collection of registered I/O handles.
     entries: Mutex<Slab<Arc<Entry>>>,
 
@@ -39,12 +42,13 @@ struct Reactor {
 
 impl Reactor {
     /// Creates a new reactor for polling I/O events.
-    fn new() -> io::Result<Reactor> {
+    pub fn new() -> io::Result<Reactor> {
         let poller = mio::Poll::new()?;
         let notify_reg = mio::Registration::new2();
 
         let mut reactor = Reactor {
             poller,
+            events: Mutex::new(mio::Events::with_capacity(1000)),
             entries: Mutex::new(Slab::new()),
             notify_reg,
             notify_token: mio::Token(0),
@@ -92,72 +96,60 @@ impl Reactor {
         Ok(())
     }
 
-    // fn notify(&self) {
-    //     self.notify_reg
-    //         .1
-    //         .set_readiness(mio::Ready::readable())
-    //         .unwrap();
-    // }
-}
+    /// Notifies the reactor so that polling stops blocking.
+    pub fn notify(&self) -> io::Result<()> {
+        self.notify_reg.1.set_readiness(mio::Ready::readable())
+    }
+
+    /// Waits on the poller for new events and wakes up tasks blocked on I/O handles.
+    ///
+    /// Returns `Ok(true)` if at least one new task was woken.
+    pub fn poll(&self, timeout: Option<Duration>) -> io::Result<bool> {
+        let mut events = self.events.lock().unwrap();
 
-/// The state of the global networking driver.
-static REACTOR: Lazy<Reactor> = Lazy::new(|| {
-    // Spawn a thread that waits on the poller for new events and wakes up tasks blocked on I/O
-    // handles.
-    std::thread::Builder::new()
-        .name("async-std/net".to_string())
-        .spawn(move || {
-            // If the driver thread panics, there's not much we can do. It is not a
-            // recoverable error and there is no place to propagate it into so we just abort.
-            abort_on_panic(|| {
-                main_loop().expect("async networking thread has panicked");
-            })
-        })
-        .expect("cannot start a thread driving blocking tasks");
-
-    Reactor::new().expect("cannot initialize reactor")
-});
-
-/// Waits on the poller for new events and wakes up tasks blocked on I/O handles.
-fn main_loop() -> io::Result<()> {
-    let reactor = &REACTOR;
-    let mut events = mio::Events::with_capacity(1000);
-
-    loop {
         // Block on the poller until at least one new event comes in.
-        reactor.poller.poll(&mut events, None)?;
+        self.poller.poll(&mut events, timeout)?;
 
         // Lock the entire entry table while we're processing new events.
-        let entries = reactor.entries.lock().unwrap();
+        let entries = self.entries.lock().unwrap();
+
+        // The number of woken tasks.
+        let mut progress = false;
 
         for event in events.iter() {
             let token = event.token();
 
-            if token == reactor.notify_token {
+            if token == self.notify_token {
                 // If this is the notification token, we just need the notification state.
-                reactor.notify_reg.1.set_readiness(mio::Ready::empty())?;
+                self.notify_reg.1.set_readiness(mio::Ready::empty())?;
             } else {
                 // Otherwise, look for the entry associated with this token.
                 if let Some(entry) = entries.get(token.0) {
                     // Set the readiness flags from this I/O event.
                     let readiness = event.readiness();
 
                     // Wake up reader tasks blocked on this I/O handle.
-                    if !(readiness & reader_interests()).is_empty() {
+                    let reader_interests = mio::Ready::all() - mio::Ready::writable();
+                    if !(readiness & reader_interests).is_empty() {
                         for w in entry.readers.lock().unwrap().drain(..) {
                             w.wake();
+                            progress = true;
                         }
                     }
 
                     // Wake up writer tasks blocked on this I/O handle.
-                    if !(readiness & writer_interests()).is_empty() {
+                    let writer_interests = mio::Ready::all() - mio::Ready::readable();
+                    if !(readiness & writer_interests).is_empty() {
                         for w in entry.writers.lock().unwrap().drain(..) {
                             w.wake();
+                            progress = true;
                         }
                     }
                 }
             }
         }
+
+        Ok(progress)
     }
 }
 
@@ -180,7 +172,8 @@ impl<T: Evented> Watcher<T> {
     /// lifetime of the returned I/O handle.
     pub fn new(source: T) -> Watcher<T> {
         Watcher {
-            entry: REACTOR
+            entry: RUNTIME
+                .reactor()
                 .register(&source)
                 .expect("cannot register an I/O event source"),
             source: Some(source),
@@ -264,7 +257,8 @@ impl<T: Evented> Watcher<T> {
     #[allow(dead_code)]
     pub fn into_inner(mut self) -> T {
         let source = self.source.take().unwrap();
-        REACTOR
+        RUNTIME
+            .reactor()
             .deregister(&source, &self.entry)
             .expect("cannot deregister I/O event source");
         source
@@ -274,7 +268,8 @@ impl<T: Evented> Watcher<T> {
 impl<T: Evented> Drop for Watcher<T> {
     fn drop(&mut self) {
         if let Some(ref source) = self.source {
-            REACTOR
+            RUNTIME
+                .reactor()
                 .deregister(source, &self.entry)
                 .expect("cannot deregister I/O event source");
         }
@@ -289,27 +284,3 @@ impl<T: Evented + fmt::Debug> fmt::Debug for Watcher<T> {
             .finish()
     }
 }
-
-/// Returns a mask containing flags that interest tasks reading from I/O handles.
-#[inline]
-fn reader_interests() -> mio::Ready {
-    mio::Ready::all() - mio::Ready::writable()
-}
-
-/// Returns a mask containing flags that interest tasks writing into I/O handles.
-#[inline]
-fn writer_interests() -> mio::Ready {
-    mio::Ready::writable() | hup()
-}
-
-/// Returns a flag containing the hangup status.
-#[inline]
-fn hup() -> mio::Ready {
-    #[cfg(unix)]
-    let ready = mio::unix::UnixReady::hup().into();
-
-    #[cfg(not(unix))]
-    let ready = mio::Ready::empty();
-
-    ready
-}