Implement simple work stealing

Author: Stjepan Glavina

Cargo.toml

@@ -28,6 +28,7 @@ unstable = []
 async-task = "1.0.0"
 cfg-if = "0.1.9"
 crossbeam-channel = "0.3.9"
+crossbeam-deque = "0.7.1"
 futures-core-preview = "0.3.0-alpha.18"
 futures-io-preview = "0.3.0-alpha.18"
 futures-timer = "0.4.0"

src/task/block_on.rs

@@ -6,11 +6,11 @@ use std::sync::Arc;
 use std::task::{RawWaker, RawWakerVTable};
 use std::thread::{self, Thread};
 
-use super::pool;
+use super::local;
 use super::task;
+use super::worker;
 use crate::future::Future;
 use crate::task::{Context, Poll, Waker};
-use crate::utils::abort_on_panic;
 
 use kv_log_macro::trace;
 
@@ -58,39 +58,36 @@ where
 
         // Log this `block_on` operation.
         let child_id = tag.task_id().as_u64();
-        let parent_id = pool::get_task(|t| t.id().as_u64()).unwrap_or(0);
+        let parent_id = worker::get_task(|t| t.id().as_u64()).unwrap_or(0);
 
         trace!("block_on", {
             parent_id: parent_id,
             child_id: child_id,
         });
 
         // Wrap the future into one that drops task-local variables on exit.
+        let future = local::add_finalizer(future);
+
         let future = async move {
             let res = future.await;
-
-            // Abort on panic because thread-local variables behave the same way.
-            abort_on_panic(|| pool::get_task(|task| task.metadata().local_map.clear()));
-
             trace!("block_on completed", {
                 parent_id: parent_id,
                 child_id: child_id,
             });
-
             res
         };
 
         // Pin the future onto the stack.
         pin_utils::pin_mut!(future);
 
-        // Transmute the future into one that is static.
+        // Transmute the future into one that is futurestatic.
         let future = mem::transmute::<
             Pin<&'_ mut dyn Future<Output = ()>>,
             Pin<&'static mut dyn Future<Output = ()>>,
         >(future);
 
         // Block on the future and and wait for it to complete.
-        pool::set_tag(&tag, || block(future));
+        worker::set_tag(&tag, || block(future));
 
         // Take out the result.
         match (*out.get()).take().unwrap() {

src/task/builder.rs

@@ -0,0 +1,32 @@
+use super::pool;
+use super::JoinHandle;
+use crate::future::Future;
+use crate::io;
+
+/// Task builder that configures the settings of a new task.
+#[derive(Debug)]
+pub struct Builder {
+    pub(crate) name: Option<String>,
+}
+
+impl Builder {
+    /// Creates a new builder.
+    pub fn new() -> Builder {
+        Builder { name: None }
+    }
+
+    /// Configures the name of the task.
+    pub fn name(mut self, name: String) -> Builder {
+        self.name = Some(name);
+        self
+    }
+
+    /// Spawns a task with the configured settings.
+    pub fn spawn<F, T>(self, future: F) -> io::Result<JoinHandle<T>>
+    where
+        F: Future<Output = T> + Send + 'static,
+        T: Send + 'static,
+    {
+        Ok(pool::get().spawn(future, self))
+    }
+}

src/task/local.rs

@@ -6,7 +6,9 @@ use std::sync::Mutex;
 
 use lazy_static::lazy_static;
 
-use super::pool;
+use super::worker;
+use crate::future::Future;
+use crate::utils::abort_on_panic;
 
 /// Declares task-local values.
 ///
@@ -152,7 +154,7 @@ impl<T: Send + 'static> LocalKey<T> {
     where
         F: FnOnce(&T) -> R,
     {
-        pool::get_task(|task| unsafe {
+        worker::get_task(|task| unsafe {
             // Prepare the numeric key, initialization function, and the map of task-locals.
             let key = self.key();
             let init = || Box::new((self.__init)()) as Box<dyn Send>;
@@ -250,3 +252,15 @@ impl Map {
         entries.clear();
     }
 }
+
+// Wrap the future into one that drops task-local variables on exit.
+pub(crate) unsafe fn add_finalizer<T>(f: impl Future<Output = T>) -> impl Future<Output = T> {
+    async move {
+        let res = f.await;
+
+        // Abort on panic because thread-local variables behave the same way.
+        abort_on_panic(|| worker::get_task(|task| task.metadata().local_map.clear()));
+
+        res
+    }
+}

src/task/mod.rs

@@ -25,15 +25,20 @@
 pub use std::task::{Context, Poll, Waker};
 
 pub use block_on::block_on;
+pub use builder::Builder;
 pub use local::{AccessError, LocalKey};
-pub use pool::{current, spawn, Builder};
+pub use pool::spawn;
 pub use sleep::sleep;
 pub use task::{JoinHandle, Task, TaskId};
+pub use worker::current;
 
 mod block_on;
+mod builder;
 mod local;
 mod pool;
 mod sleep;
+mod sleepers;
 mod task;
+mod worker;
 
 pub(crate) mod blocking;