switch to multitask

Signed-off-by: Marc-Antoine Perennou <Marc-Antoine@Perennou.com>

Author: Keruspe

Cargo.toml

@@ -29,9 +29,9 @@ default = [
   "blocking",
   "kv-log-macro",
   "log",
+  "multitask",
   "num_cpus",
   "pin-project-lite",
-  "smol",
 ]
 docs = ["attributes", "unstable", "default"]
 unstable = [
@@ -56,7 +56,7 @@ alloc = [
   "futures-core/alloc",
   "pin-project-lite",
 ]
-tokio02 = ["smol/tokio02"]
+tokio02 = ["tokio"]
 
 [dependencies]
 async-attributes = { version = "1.1.1", optional = true }
@@ -81,7 +81,7 @@ surf = { version = "1.0.3", optional = true }
 [target.'cfg(not(target_os = "unknown"))'.dependencies]
 async-io = { version = "0.1.2", optional = true }
 blocking = { version = "0.4.6", optional = true }
-smol = { version = "0.1.17", optional = true }
+multitask = { version = "0.2.0", optional = true }
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 futures-timer = { version = "3.0.2", optional = true, features = ["wasm-bindgen"] }
@@ -91,6 +91,12 @@ futures-channel = { version = "0.3.4", optional = true }
 [target.'cfg(target_arch = "wasm32")'.dev-dependencies]
 wasm-bindgen-test = "0.3.10"
 
+[dependencies.tokio]
+version = "0.2"
+default-features = false
+features = ["rt-threaded"]
+optional = true
+
 [dev-dependencies]
 femme = "1.3.0"
 rand = "0.7.3"

examples/line-count.rs

@@ -10,7 +10,7 @@ use async_std::task;
 fn main() -> io::Result<()> {
     let path = args().nth(1).expect("missing path argument");
 
-    task::block_on(async {
+    task::block_on(async move {
         let file = File::open(&path).await?;
         let mut lines = BufReader::new(file).lines();
         let mut count = 0u64;

examples/list-dir.rs

@@ -10,7 +10,7 @@ use async_std::task;
 fn main() -> io::Result<()> {
     let path = args().nth(1).expect("missing path argument");
 
-    task::block_on(async {
+    task::block_on(async move {
         let mut dir = fs::read_dir(&path).await?;
 
         while let Some(res) = dir.next().await {

examples/print-file.rs

@@ -12,7 +12,7 @@ const LEN: usize = 16 * 1024; // 16 Kb
 fn main() -> io::Result<()> {
     let path = args().nth(1).expect("missing path argument");
 
-    task::block_on(async {
+    task::block_on(async move {
         let mut file = File::open(&path).await?;
         let mut stdout = io::stdout();
         let mut buf = vec![0u8; LEN];

src/task/block_on.rs

@@ -28,7 +28,8 @@ use crate::task::Builder;
 #[cfg(not(target_os = "unknown"))]
 pub fn block_on<F, T>(future: F) -> T
 where
-    F: Future<Output = T>,
+    F: Future<Output = T> + 'static,
+    T: 'static,
 {
     Builder::new().blocking(future)
 }

src/task/builder.rs

@@ -7,10 +7,10 @@ use std::task::{Context, Poll};
 use pin_project_lite::pin_project;
 
 use crate::io;
-use crate::task::{JoinHandle, Task, TaskLocalsWrapper};
+use crate::task::{self, JoinHandle, Task, TaskLocalsWrapper};
 
 /// Task builder that configures the settings of a new task.
-#[derive(Debug, Default)]
+#[derive(Default, Debug)]
 pub struct Builder {
     pub(crate) name: Option<String>,
 }
@@ -61,9 +61,9 @@ impl Builder {
         });
 
         let task = wrapped.tag.task().clone();
-        let smol_task = smol::Task::spawn(wrapped).into();
+        let handle = task::executor::spawn(wrapped);
 
-        Ok(JoinHandle::new(smol_task, task))
+        Ok(JoinHandle::new(handle, task))
     }
 
     /// Spawns a task locally with the configured settings.
@@ -81,9 +81,9 @@ impl Builder {
         });
 
         let task = wrapped.tag.task().clone();
-        let smol_task = smol::Task::local(wrapped).into();
+        let handle = task::executor::local(wrapped);
 
-        Ok(JoinHandle::new(smol_task, task))
+        Ok(JoinHandle::new(handle, task))
     }
 
     /// Spawns a task locally with the configured settings.
@@ -141,7 +141,8 @@ impl Builder {
     #[cfg(not(target_os = "unknown"))]
     pub fn blocking<F, T>(self, future: F) -> T
     where
-        F: Future<Output = T>,
+        F: Future<Output = T> + 'static,
+        T: 'static,
     {
         let wrapped = self.build(future);
 
@@ -166,8 +167,8 @@ impl Builder {
             unsafe {
                 TaskLocalsWrapper::set_current(&wrapped.tag, || {
                     let res = if should_run {
-                        // The first call should use run.
-                        smol::run(wrapped)
+                        // The first call should run the executor
+                        task::executor::run(wrapped)
                     } else {
                         blocking::block_on(wrapped)
                     };

src/task/executor.rs

@@ -0,0 +1,92 @@
+use std::cell::RefCell;
+use std::future::Future;
+
+static GLOBAL_EXECUTOR: once_cell::sync::Lazy<multitask::Executor> = once_cell::sync::Lazy::new(multitask::Executor::new);
+
+struct Executor {
+    local_executor: multitask::LocalExecutor,
+    parker: async_io::parking::Parker,
+}
+
+thread_local! {
+    static EXECUTOR: RefCell<Executor> = RefCell::new({
+        let (parker, unparker) = async_io::parking::pair();
+        let local_executor = multitask::LocalExecutor::new(move || unparker.unpark());
+        Executor { local_executor, parker }
+    });
+}
+
+pub(crate) fn spawn<F, T>(future: F) -> multitask::Task<T>
+where
+    F: Future<Output = T> + Send + 'static,
+    T: Send + 'static,
+{
+    GLOBAL_EXECUTOR.spawn(future)
+}
+
+#[cfg(feature = "unstable")]
+pub(crate) fn local<F, T>(future: F) -> multitask::Task<T>
+where
+    F: Future<Output = T> + 'static,
+    T: 'static,
+{
+    EXECUTOR.with(|executor| executor.borrow().local_executor.spawn(future))
+}
+
+pub(crate) fn run<F, T>(future: F) -> T
+where
+    F: Future<Output = T> + 'static,
+    T: 'static,
+{
+    enter(|| EXECUTOR.with(|executor| {
+        let executor = executor.borrow();
+        let (sender, receiver) = std::sync::mpsc::channel();
+        executor.local_executor.spawn(async move {
+            sender.send(future.await).unwrap()
+        }).detach();
+        let unparker = executor.parker.unparker();
+        let global_ticker = GLOBAL_EXECUTOR.ticker(move || unparker.unpark());
+        loop {
+            let advanced = executor.local_executor.tick() || std::panic::catch_unwind(|| global_ticker.tick()).unwrap_or(true);
+            if let Ok(res) = receiver.try_recv() {
+                return res;
+            }
+            if !advanced {
+                executor.parker.park();
+            }
+        }
+    }))
+}
+
+/// Enters the tokio context if the `tokio` feature is enabled.
+fn enter<T>(f: impl FnOnce() -> T) -> T {
+    #[cfg(not(feature = "tokio02"))]
+    return f();
+
+    #[cfg(feature = "tokio02")]
+    {
+        use std::cell::Cell;
+        use tokio::runtime::Runtime;
+
+        thread_local! {
+            /// The level of nested `enter` calls we are in, to ensure that the outermost always
+            /// has a runtime spawned.
+            static NESTING: Cell<usize> = Cell::new(0);
+        }
+
+        /// The global tokio runtime.
+        static RT: once_cell::sync::Lazy<Runtime> = once_cell::sync::Lazy::new(|| Runtime::new().expect("cannot initialize tokio"));
+
+        NESTING.with(|nesting| {
+            let res = if nesting.get() == 0 {
+                nesting.replace(1);
+                RT.enter(f)
+            } else {
+                nesting.replace(nesting.get() + 1);
+                f()
+            };
+            nesting.replace(nesting.get() - 1);
+            res
+        })
+    }
+}

src/task/join_handle.rs

@@ -18,7 +18,7 @@ pub struct JoinHandle<T> {
 }
 
 #[cfg(not(target_os = "unknown"))]
-type InnerHandle<T> = async_task::JoinHandle<T, ()>;
+type InnerHandle<T> = multitask::Task<T>;
 #[cfg(target_arch = "wasm32")]
 type InnerHandle<T> = futures_channel::oneshot::Receiver<T>;
 
@@ -54,8 +54,7 @@ impl<T> JoinHandle<T> {
     #[cfg(not(target_os = "unknown"))]
     pub async fn cancel(mut self) -> Option<T> {
         let handle = self.handle.take().unwrap();
-        handle.cancel();
-        handle.await
+        handle.cancel().await
     }
 
     /// Cancel this task.
@@ -71,11 +70,6 @@ impl<T> Future for JoinHandle<T> {
     type Output = T;
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match Pin::new(&mut self.handle.as_mut().unwrap()).poll(cx) {
-            Poll::Pending => Poll::Pending,
-            Poll::Ready(output) => {
-                Poll::Ready(output.expect("cannot await the result of a panicked task"))
-            }
-        }
+        Pin::new(&mut self.handle.as_mut().unwrap()).poll(cx)
     }
 }

src/task/mod.rs

@@ -148,6 +148,8 @@ cfg_default! {
     mod block_on;
     mod builder;
     mod current;
+    #[cfg(not(target_os = "unknown"))]
+    mod executor;
     mod join_handle;
     mod sleep;
     #[cfg(not(target_os = "unknown"))]

tests/addr.rs

@@ -7,14 +7,12 @@ use async_std::task;
 
 fn blocking_resolve<A>(a: A) -> Result<Vec<SocketAddr>, String>
 where
-    A: ToSocketAddrs,
+    A: ToSocketAddrs + 'static,
     A::Iter: Send,
 {
-    let socket_addrs = task::block_on(a.to_socket_addrs());
-    match socket_addrs {
-        Ok(a) => Ok(a.collect()),
-        Err(e) => Err(e.to_string()),
-    }
+    task::block_on(async move {
+        Ok(a.to_socket_addrs().await.map_err(|e| e.to_string())?.collect())
+    })
 }
 
 #[test]
@@ -71,7 +69,7 @@ fn to_socket_addr_string() {
     let s: &str = "77.88.21.11:24352";
     assert_eq!(Ok(vec![a]), blocking_resolve(s));
 
-    let s: &String = &"77.88.21.11:24352".to_string();
+    let s: String = "77.88.21.11:24352".to_string();
     assert_eq!(Ok(vec![a]), blocking_resolve(s));
 
     let s: String = "77.88.21.11:24352".to_string();

tests/uds.rs

@@ -56,14 +56,15 @@ fn socket_ping_pong() {
     let iter_cnt = 16;
 
     let listener =
-        task::block_on(async { UnixListener::bind(&sock_path).await.expect("Socket bind") });
+        task::block_on(async move { UnixListener::bind(&sock_path).await.expect("Socket bind") });
 
     let server_handle = std::thread::spawn(move || {
-        task::block_on(async { ping_pong_server(listener, iter_cnt).await }).unwrap()
+        task::block_on(async move { ping_pong_server(listener, iter_cnt).await }).unwrap()
     });
 
+    let sock_path = tmp_dir.as_ref().join("sock");
     let client_handle = std::thread::spawn(move || {
-        task::block_on(async { ping_pong_client(&sock_path, iter_cnt).await }).unwrap()
+        task::block_on(async move { ping_pong_client(&sock_path, iter_cnt).await }).unwrap()
     });
 
     client_handle.join().unwrap();