Repository: alexcrichton/futures-rs
Pull request #665 (2017-11-30 20:43:18)

Sleep trait
Introduce a `Sleep` trait that executors can use to customize the strategy used to put the current thread to sleep. This will be used by Tokio to enable running an instance of `CurrentThread` and an instance of the Tokio reactor on a single thread.

This essentially makes the `CurrentThread` implementation generic over `W: Wakeup`. This introduces some trickiness as the thread-local state that allows spawning new futures into the scheduler cannot be generic over the `W`.

Handling this required some restructuring of `scheduler`. This work was done incrementally in each commit and each commit includes a commit message that describes the step.

/cc @aturon 

carllerche: @aturon 

I will revert the removal of `Weak`. Your assessment is correct. I had forgotten that the mpsc channel required a strong ref. Incidentally, this is why I personally don't like punting issues in complex bits of code until later. It is hard to fully reload the context.

I am worried that the weak ref will cause contention on NUMA systems, but we can punt until this is demonstrated. I will add additional comments and hopefully I won't forget the context again in the future.

Re: the `W` transmute, I think I have another completely different idea that should be safe and not require a vec. I will explore.


carllerche: @aturon the latest commit should address your feedback.
aturon: Thanks, the changes look great! Merge when ready.
carllerche: Ok, I fixed CI + updated impl comments. 

Patch:
diff --git a/benches/current_thread_executor.rs b/benches/current_thread_executor.rs
index 78da38cd..317e359a 100644
--- a/benches/current_thread_executor.rs
+++ b/benches/current_thread_executor.rs
@@ -19,7 +19,7 @@ fn execute_oneshot(b: &mut Bencher) {
     b.iter(move || {
         let cnt = Rc::new(Cell::new(0));
 
-        CurrentThread::block_with_init(|_| {
+        CurrentThread::run(|_| {
             for _ in 0..ITER {
                 let cnt = cnt.clone();
                 CurrentThread::execute(lazy(move || {
@@ -41,7 +41,7 @@ fn execute_yield_many(b: &mut Bencher) {
     b.iter(move || {
         let cnt = Rc::new(Cell::new(0));
 
-        CurrentThread::block_with_init(|_| {
+        CurrentThread::run(|_| {
             for _ in 0..TASKS {
                 let cnt = cnt.clone();
                 let mut rem = YIELDS;
@@ -83,7 +83,7 @@ fn execute_daisy(b: &mut Bencher) {
     b.iter(move || {
         cnt.set(0);
 
-        CurrentThread::block_with_init(|_| {
+        CurrentThread::run(|_| {
             daisy(DEPTH, cnt.clone());
         });
 
diff --git a/src/executor/current_thread.rs b/src/executor/current_thread.rs
index cf140210..206dd1c6 100644
--- a/src/executor/current_thread.rs
+++ b/src/executor/current_thread.rs
@@ -6,9 +6,9 @@
 //! futures, see [here].
 //!
 //! Before being able to execute futures onto [`CurrentThread`], an executor
-//! context must be setup. This is done by calling either [`block_with_init`].
-//! From within that context, [`CurrentThread::execute`] may be called with the
-//! future to run in the background.
+//! context must be setup. This is done by calling [`run`].  From within that
+//! context, [`CurrentThread::execute`] may be called with the future to run in
+//! the background.
 //!
 //! ```
 //! # use futures::executor::current_thread::*;
@@ -17,7 +17,7 @@
 //! // Calling execute here results in a panic
 //! // CurrentThread::execute(my_future);
 //!
-//! CurrentThread::block_with_init(|_| {
+//! CurrentThread::run(|_| {
 //!     // The execution context is setup, futures may be executed.
 //!     CurrentThread::execute(lazy(|| {
 //!         println!("called from the current thread executor");
@@ -28,7 +28,7 @@
 //!
 //! # Execution model
 //!
-//! When a [`CurrentThread`] execution context is setup with `block_with_init`,
+//! When a [`CurrentThread`] execution context is setup with `run`,
 //! the current thread will block and all the futures managed by the executor
 //! are driven to completion. Whenever a future receives a notification, it is
 //! pushed to the end of a scheduled list. The [`CurrentThread`] executor will
@@ -57,28 +57,27 @@
 //!
 //! [here]: https://tokio.rs/docs/going-deeper-futures/tasks/
 //! [`CurrentThread`]: struct.CurrentThread.html
-//! [`block_with_init`]: struct.CurrentThread.html#method.block_with_init
+//! [`run`]: struct.CurrentThread.html#method.run
 //! [`CurrentThread::execute`]: struct.CurrentThread.html#method.execute
 //! [`CurrentThread::execute_daemon`]: struct.CurrentThread.html#method.execute_daemon
 
 use Async;
-use executor::{self, Spawn};
+use executor::{self, Spawn, Sleep, Wakeup};
 use future::{Future, Executor, ExecuteError, ExecuteErrorKind};
 use scheduler;
 use task_impl::ThreadNotify;
 
 use std::prelude::v1::*;
 
-use std::{fmt, ptr, thread};
+use std::{fmt, thread};
 use std::cell::Cell;
 use std::rc::Rc;
-use std::sync::Arc;
 
 /// Executes futures on the current thread.
 ///
 /// All futures executed using this executor will be executed on the current
 /// thread as non-daemon futures. As such, [`CurrentThread`] will wait for these
-/// futures to complete before returning from `block_with_init`.
+/// futures to complete before returning from `run`.
 ///
 /// For more details, see the [module level](index.html) documentation.
 #[derive(Debug, Clone)]
@@ -91,7 +90,7 @@ pub struct CurrentThread {
 ///
 /// All futures executed using this executor will be executed on the current
 /// thread as daemon futures. As such, [`CurrentThread`] will **not** wait for
-/// these futures to complete before returning from `block_with_init`.
+/// these futures to complete before returning from `run`.
 ///
 /// For more details, see the [module level](index.html) documentation.
 #[derive(Debug, Clone)]
@@ -111,15 +110,14 @@ pub struct Context<'a> {
 }
 
 /// Implements the "blocking" logic for the current thread executor. A
-/// `TaskRunner` will be created during `block_with_init` and will sit on the
-/// stack until execution is complete.
+/// `TaskRunner` will be created during `run` and will sit on the stack until
+/// execution is complete.
 #[derive(Debug)]
-struct TaskRunner {
+struct TaskRunner<T> {
     /// Executes futures.
-    scheduler: Scheduler,
+    scheduler: Scheduler<T>,
 }
 
-#[derive(Debug)]
 struct CurrentRunner {
     /// When set to true, the executor should return immediately, even if there
     /// still are non-daemon futures to run.
@@ -128,13 +126,14 @@ struct CurrentRunner {
     /// Number of non-daemon futures currently being executed by the runner.
     non_daemons: Cell<usize>,
 
-    /// Raw pointer to the current scheduler.
+    /// Raw pointer to the current scheduler pusher.
     ///
     /// The raw pointer is required in order to store it in a thread-local slot.
-    scheduler: Cell<*mut Scheduler>,
+    schedule: Cell<Option<*mut Schedule>>,
 }
 
-type Scheduler = scheduler::Scheduler<SpawnedFuture, Arc<ThreadNotify>>;
+type Scheduler<T> = scheduler::Scheduler<SpawnedFuture, T>;
+type Schedule = scheduler::Schedule<SpawnedFuture>;
 
 #[derive(Debug)]
 struct SpawnedFuture {
@@ -152,7 +151,7 @@ struct Task(Spawn<Box<Future<Item = (), Error = ()>>>);
 thread_local!(static CURRENT: CurrentRunner = CurrentRunner {
     cancel: Cell::new(false),
     non_daemons: Cell::new(0),
-    scheduler: Cell::new(ptr::null_mut()),
+    schedule: Cell::new(None),
 });
 
 impl CurrentThread {
@@ -163,8 +162,7 @@ impl CurrentThread {
     /// on.
     ///
     /// The user of `CurrentThread` must ensure that when a future is submitted
-    /// to the executor, that it is done from the context of a `block_with_init`
-    /// call.
+    /// to the executor, that it is done from the context of a `run` call.
     ///
     /// For more details, see the [module level](index.html) documentation.
     pub fn current() -> CurrentThread {
@@ -180,8 +178,7 @@ impl CurrentThread {
     /// on.
     ///
     /// The user of `CurrentThread` must ensure that when a future is submitted
-    /// to the executor, that it is done from the context of a `block_with_init`
-    /// call.
+    /// to the executor, that it is done from the context of a `run` call.
     ///
     /// For more details, see the [module level](index.html) documentation.
     pub fn daemon_executor(&self) -> DaemonExecutor {
@@ -196,46 +193,59 @@ impl CurrentThread {
     /// In more detail, this function will block until:
     /// - All executing futures are complete, or
     /// - `cancel_all_executing` is invoked.
-    pub fn block_with_init<F, R>(f: F) -> R
+    pub fn run<F, R>(f: F) -> R
     where F: FnOnce(&mut Context) -> R
     {
-        TaskRunner::enter(f)
+        ThreadNotify::with_current(|mut thread_notify| {
+            TaskRunner::enter(&mut thread_notify, f)
+        })
+    }
+
+    /// Calls the given closure with a custom sleep strategy.
+    ///
+    /// This function is the same as `run` except that it allows customizing the
+    /// sleep strategy.
+    pub fn run_with_sleep<S, F, R>(sleep: &mut S, f: F) -> R
+    where F: FnOnce(&mut Context) -> R,
+          S: Sleep,
+    {
+        TaskRunner::enter(sleep, f)
     }
 
     /// Executes a future on the current thread.
     ///
     /// The provided future must complete or be canceled before
-    /// `block_with_init` will return.
+    /// `run` will return.
     ///
     /// # Panics
     ///
     /// This function can only be invoked from the context of a
-    /// `block_with_init` call; any other use will result in a panic.
+    /// `run` call; any other use will result in a panic.
     pub fn execute<F>(future: F)
     where F: Future<Item = (), Error = ()> + 'static
     {
         execute(future, false).unwrap_or_else(|_| {
             panic!("cannot call `execute` unless the thread is already \
-                    in the context of a call to `block_with_init`")
+                    in the context of a call to `run`")
         })
     }
 
     /// Executes a daemon future on the current thread.
     ///
     /// Completion of the provided future is not required for the pending
-    /// `block_with_init` call to complete. If `block_with_init` returns before
-    /// `future` completes, it will be dropped.
+    /// `run` call to complete. If `run` returns before `future` completes, it
+    /// will be dropped.
     ///
     /// # Panics
     ///
     /// This function can only be invoked from the context of a
-    /// `block_with_init` call; any other use will result in a panic.
+    /// `run` call; any other use will result in a panic.
     pub fn execute_daemon<F>(future: F)
     where F: Future<Item = (), Error = ()> + 'static
     {
         execute(future, true).unwrap_or_else(|_| {
             panic!("cannot call `execute` unless the thread is already \
-                    in the context of a call to `block_with_init`")
+                    in the context of a call to `run`")
         })
     }
 
@@ -246,12 +256,12 @@ impl CurrentThread {
     /// # Panics
     ///
     /// This function can only be invoked from the context of a
-    /// `block_with_init` call; any other use will result in a panic.
+    /// `run` call; any other use will result in a panic.
     pub fn cancel_all_executing() {
         CurrentRunner::with(|runner| runner.cancel_all_executing())
             .unwrap_or_else(|()| {
                 panic!("cannot call `cancel_all_executing` unless the thread is already \
-                        in the context of a call to `block_with_init`")
+                        in the context of a call to `run`")
             })
     }
 }
@@ -284,40 +294,45 @@ impl<'a> Context<'a> {
 /// checking the thread-local variable tracking the current executor.
 ///
 /// If this function is not called in context of an executor, i.e. outside of
-/// `block_with_init`, then `Err` is returned.
+/// `run`, then `Err` is returned.
 ///
 /// This function does not panic.
 fn execute<F>(future: F, daemon: bool) -> Result<(), ExecuteError<F>>
 where F: Future<Item = (), Error = ()> + 'static,
 {
     CURRENT.with(|current| {
-        if current.scheduler.get().is_null() {
-            Err(ExecuteError::new(ExecuteErrorKind::Shutdown, future))
-        } else {
-            let spawned = SpawnedFuture {
-                daemon: daemon,
-                inner: Task::new(future),
-            };
+        match current.schedule.get() {
+            Some(schedule) => {
+                let spawned = SpawnedFuture {
+                    daemon: daemon,
+                    inner: Task::new(future),
+                };
 
-            if !daemon {
-                let non_daemons = current.non_daemons.get();
-                current.non_daemons.set(non_daemons + 1);
-            }
+                if !daemon {
+                    let non_daemons = current.non_daemons.get();
+                    current.non_daemons.set(non_daemons + 1);
+                }
 
-            unsafe {
-                (*current.scheduler.get()).push(spawned);
-            }
+                unsafe { (*schedule).schedule(spawned); }
 
-            Ok(())
+                Ok(())
+            }
+            None => {
+                Err(ExecuteError::new(ExecuteErrorKind::Shutdown, future))
+            }
         }
     })
 }
 
-impl TaskRunner {
+impl<T> TaskRunner<T>
+where T: Wakeup,
+{
     /// Return a new `TaskRunner`
-    fn new(thread_notify: Arc<ThreadNotify>) -> TaskRunner {
+    fn new(wakeup: T) -> TaskRunner<T> {
+        let scheduler = scheduler::Scheduler::new(wakeup);
+
         TaskRunner {
-            scheduler: scheduler::Scheduler::new(thread_notify),
+            scheduler: scheduler,
         }
     }
 
@@ -332,9 +347,9 @@ impl TaskRunner {
     ///
     /// Before any future is polled, the scheduler must be set to a thread-local
     /// variable so that `execute` is able to submit new futures to the current
-    /// executor. Because `Scheduler::push` requires `&mut self`, this
+    /// executor. Because `Scheduler::schedule` requires `&mut self`, this
     /// introduces a mutability hazard. This hazard is minimized with some
-    /// indirection. See `set_scheduler` for more details.
+    /// indirection. See `set_schedule` for more details.
     ///
     /// Once all context is setup, the init closure is invoked. This is the
     /// "boostrapping" process that executes the initial futures into the
@@ -342,57 +357,57 @@ impl TaskRunner {
     /// state until all non daemon futures complete. When no scheduled futures
     /// are ready to be advanced, the thread is blocked using
     /// `ThreadNotify::park`.
-    fn enter<F, R>(f: F) -> R
+    fn enter<S, F, R>(sleep: &mut S, f: F) -> R
     where F: FnOnce(&mut Context) -> R,
+          S: Sleep<Wakeup = T>,
     {
-        // Create a new task runner that will be used for the duration of `f`.
-        ThreadNotify::with_current(|thread_notify| {
-            let mut runner = TaskRunner::new(thread_notify.clone());
+        let mut runner = TaskRunner::new(sleep.wakeup());
 
-            CURRENT.with(|current| {
-                // Make sure that another task runner is not set.
-                //
-                // This should not be ever possible due to how `set_scheduler`
-                // is setup, but better safe than sorry!
-                assert!(current.scheduler.get().is_null());
-
-                let enter = executor::enter()
-                    .expect("cannot execute `CurrentThread` executor from within \
-                             another executor");
-
-                // Enter an execution scope
-                let mut ctx = Context {
-                    enter: enter,
-                    _p: ::std::marker::PhantomData,
-                };
+        CURRENT.with(|current| {
+            // Make sure that another task runner is not set.
+            //
+            // This should not be ever possible due to how `set_schedule`
+            // is setup, but better safe than sorry!
+            assert!(current.schedule.get().is_none());
+
+            let enter = executor::enter()
+                .expect("cannot execute `CurrentThread` executor from within \
+                         another executor");
+
+            // Enter an execution scope
+            let mut ctx = Context {
+                enter: enter,
+                _p: ::std::marker::PhantomData,
+            };
 
-                // Set the scheduler to the TLS and perform setup work,
-                // returning a future to execute.
-                //
-                // This could possibly suubmit other futures for execution.
-                let ret = current.set_scheduler(&mut runner.scheduler, || {
-                    f(&mut ctx)
-                });
+            // Set the scheduler to the TLS and perform setup work,
+            // returning a future to execute.
+            //
+            // This could possibly suubmit other futures for execution.
+            let ret = current.set_schedule(&mut runner.scheduler as &mut Schedule, || {
+                f(&mut ctx)
+            });
 
-                // Execute the runner.
-                //
-                // This function will not return until either
-                //
-                // a) All non daemon futures have completed execution
-                // b) `cancel_all_executing` is called, forcing the executor to
-                // return.
-                runner.run(thread_notify, current);
+            // Execute the runner.
+            //
+            // This function will not return until either
+            //
+            // a) All non daemon futures have completed execution
+            // b) `cancel_all_executing` is called, forcing the executor to
+            // return.
+            runner.run(sleep, current);
 
-                // Not technically required, but this makes the fact that `ctx`
-                // needs to live until this point explicit.
-                drop(ctx);
+            // Not technically required, but this makes the fact that `ctx`
+            // needs to live until this point explicit.
+            drop(ctx);
 
-                ret
-            })
+            ret
         })
     }
 
-    fn run(&mut self, thread_notify: &Arc<ThreadNotify>, current: &CurrentRunner) {
+    fn run<S>(&mut self, sleep: &mut S, current: &CurrentRunner)
+    where S: Sleep<Wakeup = T>,
+    {
         use scheduler::Tick;
 
         while current.is_running() {
@@ -404,13 +419,13 @@ impl TaskRunner {
                 //
                 // This lets us push new futures into the scheduler. It also
                 // lets us pass the scheduler mutable reference into
-                // `set_scheduler`, which sets the thread-local variable that
+                // `set_schedule`, which sets the thread-local variable that
                 // `CurrentThread::execute` uses for submitting new futures to the
                 // "current" executor.
                 //
-                // See `set_scheduler` documentation for more details on how we
+                // See `set_schedule` documentation for more details on how we
                 // guard against mutable pointer aliasing.
-                current.set_scheduler(scheduler, || {
+                current.set_schedule(scheduler as &mut Schedule, || {
                     match spawned.inner.0.poll_future_notify(notify, 0) {
                         Ok(Async::Ready(_)) | Err(_) => {
                             Async::Ready(spawned.daemon)
@@ -442,7 +457,7 @@ impl TaskRunner {
 
                     // Block the current thread until a future managed by the scheduler
                     // receives a readiness notification.
-                    thread_notify.park();
+                    sleep.sleep();
                 }
                 Tick::Inconsistent => {
                     // Yield the thread and loop
@@ -458,15 +473,15 @@ impl CurrentRunner {
     where F: FnOnce(&Self) -> R,
     {
         CURRENT.with(|current| {
-            if current.scheduler.get().is_null() {
-                Err(())
-            } else {
+            if current.schedule.get().is_some() {
                 Ok(f(current))
+            } else {
+                Err(())
             }
         })
     }
 
-    /// Set the provided scheduler to the TLS slot for the duration of the
+    /// Set the provided schedule handle to the TLS slot for the duration of the
     /// closure.
     ///
     /// `CurrentThread::execute` will access the CURRENT thread-local variable in
@@ -479,16 +494,16 @@ impl CurrentRunner {
     /// that one can not access the scheduler anywhere else.
     ///
     /// To do this, we only allow access to the thread local variable from
-    /// within the closure passed to `set_scheduler`. This function also takes a
+    /// within the closure passed to `set_schedule`. This function also takes a
     /// &mut reference to the scheduler, which is essentially holding a "lock"
     /// on that reference, preventing any other location in the code from
     /// also getting that &mut reference.
     ///
-    /// When `set_scheduler` returns, the thread-local variable containing the
+    /// When `set_schedule` returns, the thread-local variable containing the
     /// mut reference is set to null. This is done even if the closure panics.
     ///
     /// This reduces the odds of introducing pointer aliasing.
-    fn set_scheduler<F, R>(&self, scheduler: &mut Scheduler, f: F) -> R
+    fn set_schedule<F, R>(&self, schedule: &mut Schedule, f: F) -> R
     where F: FnOnce() -> R
     {
         // Ensure that the runner is removed from the thread-local context
@@ -497,13 +512,13 @@ impl CurrentRunner {
 
         impl<'a> Drop for Reset<'a> {
             fn drop(&mut self) {
-                self.0.scheduler.set(ptr::null_mut());
+                self.0.schedule.set(None);
             }
         }
 
         let _reset = Reset(self);
 
-        self.scheduler.set(scheduler as *mut Scheduler);
+        self.schedule.set(Some(schedule as *mut Schedule));
 
         f()
     }
diff --git a/src/executor/mod.rs b/src/executor/mod.rs
index c8e0643e..957a0ca7 100644
--- a/src/executor/mod.rs
+++ b/src/executor/mod.rs
@@ -13,6 +13,9 @@ pub mod current_thread;
 #[cfg(feature = "use_std")]
 mod enter;
 
+#[cfg(feature = "use_std")]
+mod sleep;
+
 #[allow(deprecated)]
 #[cfg(feature = "use_std")]
 pub use task_impl::{Unpark, Executor, Run};
@@ -26,3 +29,6 @@ pub use self::current_thread::CurrentThread;
 
 #[cfg(feature = "use_std")]
 pub use self::enter::{enter, Enter};
+
+#[cfg(feature = "use_std")]
+pub use self::sleep::{Sleep, Wakeup};
diff --git a/src/executor/sleep.rs b/src/executor/sleep.rs
new file mode 100644
index 00000000..dbdaa8db
--- /dev/null
+++ b/src/executor/sleep.rs
@@ -0,0 +1,22 @@
+use std::time::Duration;
+
+/// Puts the current thread to sleep.
+pub trait Sleep {
+    /// Wake up handle.
+    type Wakeup: Wakeup;
+
+    /// Get a new `Wakeup` handle.
+    fn wakeup(&self) -> Self::Wakeup;
+
+    /// Put the current thread to sleep.
+    fn sleep(&mut self);
+
+    /// Put the current thread to sleep for at most `duration`.
+    fn sleep_timeout(&mut self, duration: Duration);
+}
+
+/// Wake up a sleeping thread.
+pub trait Wakeup: Send + Sync + 'static {
+    /// Wake up the sleeping thread.
+    fn wakeup(&self);
+}
diff --git a/src/scheduler.rs b/src/scheduler.rs
index f343b33b..1db32d31 100644
--- a/src/scheduler.rs
+++ b/src/scheduler.rs
@@ -1,8 +1,7 @@
 //! An unbounded set of futures.
 
 use Async;
-use executor::{self, UnsafeNotify, NotifyHandle};
-use task_impl::{AtomicTask, ThreadNotify};
+use executor::{self, UnsafeNotify, NotifyHandle, Wakeup};
 
 use std::cell::UnsafeCell;
 use std::fmt::{self, Debug};
@@ -19,45 +18,55 @@ use std::usize;
 /// This is used both by `FuturesUnordered` and the current-thread executor.
 pub struct Scheduler<T, W> {
     inner: Arc<Inner<T, W>>,
-    len: usize,
-    head_all: *const Node<T, W>,
+    nodes: List<T, W>,
+}
+
+/// Schedule new futures
+pub trait Schedule<T> {
+    /// Schedule a new future.
+    fn schedule(&mut self, item: T);
 }
 
 pub struct Notify<'a, T: 'a, W: 'a>(&'a Arc<Node<T, W>>);
 
-/// Wakeup a sleeper
-///
-/// This notifies a sleeper that the the `Scheduler` is ready to do more work.
-pub trait Wakeup: Send + Sync {
-    /// Wakeup a sleeper
-    fn wakeup(&self);
+// A linked-list of nodes
+struct List<T, W> {
+    len: usize,
+    head: *const Node<T, W>,
+    tail: *const Node<T, W>,
 }
 
 unsafe impl<T: Send, W: Wakeup> Send for Scheduler<T, W> {}
 unsafe impl<T: Sync, W: Wakeup> Sync for Scheduler<T, W> {}
 
-// Scheduler is implemented using two linked lists. One which links all
-// items managed by a `Scheduler` and one that tracks items that have
-// been scheduled for polling. The first linked list is not thread safe and is
-// only accessed by the thread that owns the `Scheduler` value. The
-// second linked list is an implementation of the intrusive MPSC queue algorithm
-// described by 1024cores.net.
+// Scheduler is implemented using two linked lists. The first linked list tracks
+// all items managed by a `Scheduler`. This list is stored on the `Scheduler`
+// struct and is **not** thread safe. The second linked list is an
+// implementation of the intrusive MPSC queue algorithm described by
+// 1024cores.net and is stored on `Inner`. This linked list can push items to
+// the back concurrently but only one consumer may pop from the front. To
+// enforce this requirement, all popping will be performed via fns on
+// `Scheduler` that take `&mut self`.
 //
 // When a item is submitted to the set a node is allocated and inserted in
-// both linked lists. The next call to `poll` will (eventually) see this node
-// and call `poll` on the item.
+// both linked lists. This means that all insertion operations **must** be
+// originated from `Scheduler` with `&mut self` The next call to `tick` will
+// (eventually) see this node and call `poll` on the item.
 //
-// Before a managed item is polled, the current task's `Notify` is replaced
-// with one that is aware of the specific item being run. This ensures that
-// task notifications generated by that specific item are visible to
-// `Scheduler`. When a notification is received, the node is scheduled
-// for polling by being inserted into the concurrent linked list.
+// Nodes are wrapped in `Arc` cells which manage the lifetime of the node.
+// However, `Arc` handles are sometimes cast to `*const Node` pointers.
+// Specifically, when a node is stored in at least one of the two lists
+// described above, this represents a logical `Arc` handle. This is how
+// `Scheduler` maintains its reference to all nodes it manages. Each
+// `NotifyHande` instance is an `Arc<Node>` as well.
 //
-// Each node uses an `AtomicUsize` to track it's state. The node state is the
-// reference count (the number of outstanding handles to the node) as well as a
-// flag tracking if the node is currently inserted in the atomic queue. When the
-// item is notified, it will only insert itself into the linked list if it
-// isn't currently inserted.
+// When `Scheduler` drops, it clears the linked list of all nodes that it
+// manages. When doing so, it must attempt to decrement the reference count (by
+// dropping an Arc handle). However, it can **only** decrement the reference
+// count if the node is not currently stored in the mpsc channel. If the node
+// **is** "queued" in the mpsc channel, then the arc reference count cannot be
+// decremented. Once the node is popped from the mpsc channel, then the final
+// arc reference count can be decremented, thus freeing the node.
 
 #[allow(missing_debug_implementations)]
 struct Inner<T, W> {
@@ -85,11 +94,11 @@ struct Node<T, W> {
     // Next pointer in readiness queue
     next_readiness: AtomicPtr<Node<T, W>>,
 
-    // Queue that we'll be enqueued to when notified
-    queue: Weak<Inner<T, W>>,
-
     // Whether or not this node is currently in the mpsc queue.
     queued: AtomicBool,
+
+    // Queue that we'll be enqueued to when notified
+    queue: Weak<Inner<T, W>>,
 }
 
 /// Returned by the `Scheduler::tick` function, allowing the caller to decide
@@ -139,9 +148,8 @@ where W: Wakeup,
         });
 
         Scheduler {
-            len: 0,
-            head_all: ptr::null_mut(),
             inner: inner,
+            nodes: List::new(),
         }
     }
 }
@@ -156,40 +164,12 @@ impl<T, W: Wakeup> Scheduler<T, W> {
     ///
     /// This represents the total number of in-flight items.
     pub fn len(&self) -> usize {
-        self.len
+        self.nodes.len
     }
 
     /// Returns `true` if the set contains no items
     pub fn is_empty(&self) -> bool {
-        self.len == 0
-    }
-
-    /// Push a item into the set.
-    ///
-    /// This function submits the given item to the set for managing. This
-    /// function will not call `poll` on the submitted item. The caller must
-    /// ensure that `Scheduler::poll` is called in order to receive task
-    /// notifications.
-    pub fn push(&mut self, item: T) {
-        let node = Arc::new(Node {
-            item: UnsafeCell::new(Some(item)),
-            next_all: UnsafeCell::new(ptr::null_mut()),
-            prev_all: UnsafeCell::new(ptr::null_mut()),
-            next_readiness: AtomicPtr::new(ptr::null_mut()),
-            queued: AtomicBool::new(true),
-            queue: Arc::downgrade(&self.inner),
-        });
-
-        // Right now our node has a strong reference count of 1. We transfer
-        // ownership of this reference count to our internal linked list
-        // and we'll reclaim ownership through the `unlink` function below.
-        let ptr = self.link(node);
-
-        // We'll need to get the item "into the system" to start tracking it,
-        // e.g. getting its unpark notifications going to us tracking which
-        // items are ready. To do that we unconditionally enqueue it for
-        // polling here.
-        self.inner.enqueue(ptr);
+        self.nodes.len == 0
     }
 
     /// Advance the scheduler state.
@@ -213,28 +193,18 @@ impl<T, W: Wakeup> Scheduler<T, W> {
             debug_assert!(node != self.inner.stub());
 
             unsafe {
-                let mut item = match (*(*node).item.get()).take() {
-                    Some(item) => item,
-
-                    // If the item has already gone away then we're just
-                    // cleaning out this node. See the comment in
-                    // `release_node` for more information, but we're basically
-                    // just taking ownership of our reference count here.
-                    None => {
-                        let node = ptr2arc(node);
-                        assert!((*node.next_all.get()).is_null());
-                        assert!((*node.prev_all.get()).is_null());
-                        continue
-                    }
+                if (*(*node).item.get()).is_none() {
+                    // The node has already been released. However, while it was
+                    // being released, another thread notified it, which
+                    // resulted in it getting pushed into the mpsc channel.
+                    //
+                    // In this case, we just dec the ref count.
+                    let node = ptr2arc(node);
+                    assert!((*node.next_all.get()).is_null());
+                    assert!((*node.prev_all.get()).is_null());
+                    continue
                 };
 
-                // Unset queued flag... this must be done before
-                // polling. This ensures that the item gets
-                // rescheduled if it is notified **during** a call
-                // to `poll`.
-                let prev = (*node).queued.swap(false, SeqCst);
-                assert!(prev);
-
                 // We're going to need to be very careful if the `poll`
                 // function below panics. We need to (a) not leak memory and
                 // (b) ensure that we still don't have any use-after-frees. To
@@ -243,54 +213,75 @@ impl<T, W: Wakeup> Scheduler<T, W> {
                 // * This "bomb" here will call `release_node` if dropped
                 //   abnormally. That way we'll be sure the memory management
                 //   of the `node` is managed correctly.
-                // * The item was extracted above (taken ownership). That way
-                //   if it panics we're guaranteed that the item is
-                //   dropped on this thread and doesn't accidentally get
-                //   dropped on a different thread (bad).
+                //
                 // * We unlink the node from our internal queue to preemptively
-                //   assume it'll panic, in which case we'll want to discard it
-                //   regardless.
+                //   assume is is complete (will return Ready or panic), in
+                //   which case we'll want to discard it regardless.
+                //
                 struct Bomb<'a, T: 'a, W: 'a> {
                     queue: &'a mut Scheduler<T, W>,
                     node: Option<Arc<Node<T, W>>>,
                 }
+
                 impl<'a, T, W> Drop for Bomb<'a, T, W> {
                     fn drop(&mut self) {
                         if let Some(node) = self.node.take() {
-                            self.queue.release_node(node);
+                            release_node(node);
                         }
                     }
                 }
+
                 let mut bomb = Bomb {
-                    node: Some(self.unlink(node)),
+                    node: Some(self.nodes.remove(node)),
                     queue: self,
                 };
 
-                // Poll the underlying item with the appropriate `notify`
-                // implementation. This is where a large bit of the unsafety
-                // starts to stem from internally. The `notify` instance itself
-                // is basically just our `Arc<Node<T>>` and tracks the mpsc
-                // queue of ready items.
-                //
-                // Critically though `Node<T>` won't actually access `T`, the
-                // item, while it's floating around inside of `Task`
-                // instances. These structs will basically just use `T` to size
-                // the internal allocation, appropriately accessing fields and
-                // deallocating the node if need be.
+                // Now that the bomb holds the node, create a new scope. This
+                // scope ensures that the borrow will go out of scope before we
+                // mutate the node pointer in `bomb` again
                 let res = {
+                    let node = bomb.node.as_ref().unwrap();
+
+                    // Get a reference to the inner future. We already ensured
+                    // that the item `is_some`.
+                    let item = (*node.item.get()).as_mut().unwrap();
+
+                    // Unset queued flag... this must be done before
+                    // polling. This ensures that the item gets
+                    // rescheduled if it is notified **during** a call
+                    // to `poll`.
+                    let prev = (*node).queued.swap(false, SeqCst);
+                    assert!(prev);
+
+                    // Poll the underlying item with the appropriate `notify`
+                    // implementation. This is where a large bit of the unsafety
+                    // starts to stem from internally. The `notify` instance itself
+                    // is basically just our `Arc<Node<T>>` and tracks the mpsc
+                    // queue of ready items.
+                    //
+                    // Critically though `Node<T>` won't actually access `T`, the
+                    // item, while it's floating around inside of `Task`
+                    // instances. These structs will basically just use `T` to size
+                    // the internal allocation, appropriately accessing fields and
+                    // deallocating the node if need be.
                     let queue = &mut *bomb.queue;
                     let notify = Notify(bomb.node.as_ref().unwrap());
-                    f(queue, &mut item, &notify)
+                    f(queue, item, &notify)
                 };
 
                 let ret = match res {
                     Async::NotReady => {
+                        // The future is not done, push it back into the "all
+                        // node" list.
                         let node = bomb.node.take().unwrap();
-                        *node.item.get() = Some(item);
-                        bomb.queue.link(node);
+                        bomb.queue.nodes.push_back(node);
                         continue;
                     }
-                    Async::Ready(v) => Tick::Data(v),
+                    Async::Ready(v) => {
+                        // `bomb` will take care of unlinking and releasing the
+                        // node.
+                        Tick::Data(v)
+                    }
                 };
 
                 return ret
@@ -300,78 +291,59 @@ impl<T, W: Wakeup> Scheduler<T, W> {
 
     /// Returns an iterator that allows modifying each item in the set.
     pub fn iter_mut(&mut self) -> IterMut<T, W> {
-        IterMut {
-            node: self.head_all,
-            len: self.len,
-            _marker: PhantomData
-        }
+        self.nodes.iter_mut()
     }
 }
 
-impl<T, W> Scheduler<T, W> {
-    fn release_node(&mut self, node: Arc<Node<T, W>>) {
-        // The item is done, try to reset the queued flag. This will prevent
-        // `notify` from doing any work in the item
-        let prev = node.queued.swap(true, SeqCst);
+impl<T, W: Wakeup> Schedule<T> for Scheduler<T, W> {
+    fn schedule(&mut self, item: T) {
+        let node = Arc::new(Node {
+            item: UnsafeCell::new(Some(item)),
+            next_all: UnsafeCell::new(ptr::null_mut()),
+            prev_all: UnsafeCell::new(ptr::null_mut()),
+            next_readiness: AtomicPtr::new(ptr::null_mut()),
+            queued: AtomicBool::new(true),
+            queue: Arc::downgrade(&self.inner),
+        });
 
-        // Drop the item, even if it hasn't finished yet. This is safe
-        // because we're dropping the item on the thread that owns
-        // `Scheduler`, which correctly tracks T's lifetimes and such.
-        unsafe {
-            drop((*node.item.get()).take());
-        }
+        // Right now our node has a strong reference count of 1. We transfer
+        // ownership of this reference count to our internal linked list
+        // and we'll reclaim ownership through the `unlink` function below.
+        let ptr = self.nodes.push_back(node);
 
-        // If the queued flag was previously set then it means that this node
-        // is still in our internal mpsc queue. We then transfer ownership
-        // of our reference count to the mpsc queue, and it'll come along and
-        // free it later, noticing that the item is `None`.
-        //
-        // If, however, the queued flag was *not* set then we're safe to
-        // release our reference count on the internal node. The queued flag
-        // was set above so all item `enqueue` operations will not actually
-        // enqueue the node, so our node will never see the mpsc queue again.
-        // The node itself will be deallocated once all reference counts have
-        // been dropped by the various owning tasks elsewhere.
-        if prev {
-            mem::forget(node);
-        }
+        // We'll need to get the item "into the system" to start tracking it,
+        // e.g. getting its unpark notifications going to us tracking which
+        // items are ready. To do that we unconditionally enqueue it for
+        // polling here.
+        self.inner.enqueue(ptr);
     }
+}
 
-    /// Insert a new node into the internal linked list.
-    fn link(&mut self, node: Arc<Node<T, W>>) -> *const Node<T, W> {
-        let ptr = arc2ptr(node);
-        unsafe {
-            *(*ptr).next_all.get() = self.head_all;
-            if !self.head_all.is_null() {
-                *(*self.head_all).prev_all.get() = ptr;
-            }
-        }
+fn release_node<T, W>(node: Arc<Node<T, W>>) {
+    // The item is done, try to reset the queued flag. This will prevent
+    // `notify` from doing any work in the item
+    let prev = node.queued.swap(true, SeqCst);
 
-        self.head_all = ptr;
-        self.len += 1;
-        return ptr
+    // Drop the item, even if it hasn't finished yet. This is safe
+    // because we're dropping the item on the thread that owns
+    // `Scheduler`, which correctly tracks T's lifetimes and such.
+    unsafe {
+        drop((*node.item.get()).take());
     }
 
-    /// Remove the node from the linked list tracking all nodes currently
-    /// managed by `Scheduler`.
-    unsafe fn unlink(&mut self, node: *const Node<T, W>) -> Arc<Node<T, W>> {
-        let node = ptr2arc(node);
-        let next = *node.next_all.get();
-        let prev = *node.prev_all.get();
-        *node.next_all.get() = ptr::null_mut();
-        *node.prev_all.get() = ptr::null_mut();
-
-        if !next.is_null() {
-            *(*next).prev_all.get() = prev;
-        }
-
-        if !prev.is_null() {
-            *(*prev).next_all.get() = next;
-        } else {
-            self.head_all = next;
-        }
-        self.len -= 1;
-        return node
+    // If the queued flag was previously set then it means that this node
+    // is still in our internal mpsc queue. We then transfer ownership
+    // of our reference count to the mpsc queue, and it'll come along and
+    // free it later, noticing that the item is `None`.
+    //
+    // If, however, the queued flag was *not* set then we're safe to
+    // release our reference count on the internal node. The queued flag
+    // was set above so all item `enqueue` operations will not actually
+    // enqueue the node, so our node will never see the mpsc queue again.
+    // The node itself will be deallocated once all reference counts have
+    // been dropped by the various owning tasks elsewhere.
+    if prev {
+        mem::forget(node);
     }
 }
 
@@ -388,12 +360,8 @@ impl<T, W> Drop for Scheduler<T, W> {
         // flying around which contain `Node<T>` references inside them. We'll
         // let those naturally get deallocated when the `Task` itself goes out
         // of scope or gets notified.
-        unsafe {
-            while !self.head_all.is_null() {
-                let head = self.head_all;
-                let node = self.unlink(head);
-                self.release_node(node);
-            }
+        while let Some(node) = self.nodes.pop_front() {
+            release_node(node);
         }
 
         // Note that at this point we could still have a bunch of nodes in the
@@ -524,16 +492,97 @@ impl<T, W> Drop for Inner<T, W> {
     }
 }
 
-impl Wakeup for AtomicTask {
-    fn wakeup(&self) {
-        self.notify();
+impl<T, W> List<T, W> {
+    fn new() -> Self {
+        List {
+            len: 0,
+            head: ptr::null_mut(),
+            tail: ptr::null_mut(),
+        }
     }
-}
 
-impl Wakeup for Arc<ThreadNotify> {
-    fn wakeup(&self) {
-        use executor::Notify;
-        self.notify(0);
+    /// Prepends an element to the back of the list
+    fn push_back(&mut self, node: Arc<Node<T, W>>) -> *const Node<T, W> {
+        let ptr = arc2ptr(node);
+
+        unsafe {
+            // Point to the current last node in the list
+            *(*ptr).prev_all.get() = self.tail;
+            *(*ptr).next_all.get() = ptr::null_mut();
+
+            if !self.tail.is_null() {
+                *(*self.tail).next_all.get() = ptr;
+                self.tail = ptr;
+            } else {
+                // This is the first node
+                self.tail = ptr;
+                self.head = ptr;
+            }
+        }
+
+        self.len += 1;
+
+        return ptr
+    }
+
+    /// Pop an element from the front of the list
+    fn pop_front(&mut self) -> Option<Arc<Node<T, W>>> {
+        if self.head.is_null() {
+            // The list is empty
+            return None;
+        }
+
+        self.len -= 1;
+
+        unsafe {
+            // Convert the ptr to Arc<_>
+            let node = ptr2arc(self.head);
+
+            // Update the head pointer
+            self.head = *node.next_all.get();
+
+            // If the pointer is null, then the list is empty
+            if self.head.is_null() {
+                self.tail = ptr::null_mut();
+            } else {
+                *(*self.head).prev_all.get() = ptr::null_mut();
+            }
+
+            Some(node)
+        }
+    }
+
+    /// Remove a specific node
+    unsafe fn remove(&mut self, node: *const Node<T, W>) -> Arc<Node<T, W>> {
+        let node = ptr2arc(node);
+        let next = *node.next_all.get();
+        let prev = *node.prev_all.get();
+        *node.next_all.get() = ptr::null_mut();
+        *node.prev_all.get() = ptr::null_mut();
+
+        if !next.is_null() {
+            *(*next).prev_all.get() = prev;
+        } else {
+            self.tail = prev;
+        }
+
+        if !prev.is_null() {
+            *(*prev).next_all.get() = next;
+        } else {
+            self.head = next;
+        }
+
+        self.len -= 1;
+
+        return node
+    }
+
+    fn iter_mut(&mut self) -> IterMut<T, W> {
+        IterMut {
+            node: self.head,
+            len: self.len,
+            _marker: PhantomData
+        }
     }
 }
 
diff --git a/src/stream/futures_unordered.rs b/src/stream/futures_unordered.rs
index 4d087545..313bd206 100644
--- a/src/stream/futures_unordered.rs
+++ b/src/stream/futures_unordered.rs
@@ -3,7 +3,7 @@
 use std::iter::FromIterator;
 
 use {task, Stream, Future, Poll, Async};
-use scheduler::{self, Scheduler};
+use scheduler::{self, Scheduler, Schedule};
 use task_impl::{self, AtomicTask};
 
 /// An unbounded set of futures.
@@ -73,7 +73,7 @@ impl<T> FuturesUnordered<T> {
     /// ensure that `FuturesUnordered::poll` is called in order to receive task
     /// notifications.
     pub fn push(&mut self, future: T) {
-        self.inner.push(future)
+        self.inner.schedule(future)
     }
 
     /// Returns an iterator that allows modifying each future in the set.
diff --git a/src/task_impl/atomic_task.rs b/src/task_impl/atomic_task.rs
index aa368273..029e108d 100644
--- a/src/task_impl/atomic_task.rs
+++ b/src/task_impl/atomic_task.rs
@@ -175,6 +175,13 @@ impl AtomicTask {
     }
 }
 
+#[cfg(feature = "use_std")]
+impl ::executor::Wakeup for AtomicTask {
+    fn wakeup(&self) {
+        self.notify();
+    }
+}
+
 impl fmt::Debug for AtomicTask {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         write!(fmt, "AtomicTask")
diff --git a/src/task_impl/std/mod.rs b/src/task_impl/std/mod.rs
index 04efb58c..e37e4655 100644
--- a/src/task_impl/std/mod.rs
+++ b/src/task_impl/std/mod.rs
@@ -7,6 +7,7 @@ use std::mem;
 use std::ptr;
 use std::sync::{Arc, Mutex, Condvar, Once, ONCE_INIT};
 use std::sync::atomic::{AtomicUsize, Ordering};
+use std::time::{Duration, Instant};
 
 use {Future, Stream, Sink, Poll, Async, StartSend, AsyncSink};
 use super::core;
@@ -500,6 +501,10 @@ impl ThreadNotify {
     }
 
     pub fn park(&self) {
+        self.park_timeout(None);
+    }
+
+    pub fn park_timeout(&self, dur: Option<Duration>) {
         // If currently notified, then we skip sleeping. This is checked outside
         // of the lock to avoid acquiring a mutex if not necessary.
         match self.state.compare_and_swap(NOTIFY, IDLE, Ordering::SeqCst) {
@@ -524,9 +529,26 @@ impl ThreadNotify {
             _ => unreachable!(),
         }
 
-        // Loop until we've been notified
+        // Track (until, remaining)
+        let mut time = dur.map(|dur| (Instant::now() + dur, dur));
+
         loop {
-            m = self.condvar.wait(m).unwrap();
+            m = match time {
+                Some((until, rem)) => {
+                    let (guard, _) = self.condvar.wait_timeout(m, rem).unwrap();
+                    let now = Instant::now();
+
+                    if now >= until {
+                        // Timed out... exit sleep state
+                        self.state.store(IDLE, Ordering::SeqCst);
+                        return;
+                    }
+
+                    time = Some((until, until - now));
+                    guard
+                }
+                None => self.condvar.wait(m).unwrap(),
+            };
 
             // Transition back to idle, loop otherwise
             if NOTIFY == self.state.compare_and_swap(NOTIFY, IDLE, Ordering::SeqCst) {
@@ -534,10 +556,8 @@ impl ThreadNotify {
             }
         }
     }
-}
 
-impl Notify for ThreadNotify {
-    fn notify(&self, _unpark_id: usize) {
+    fn unpark(&self) {
         // First, try transitioning from IDLE -> NOTIFY, this does not require a
         // lock.
         match self.state.compare_and_swap(IDLE, NOTIFY, Ordering::SeqCst) {
@@ -560,6 +580,34 @@ impl Notify for ThreadNotify {
     }
 }
 
+impl Notify for ThreadNotify {
+    fn notify(&self, _unpark_id: usize) {
+        self.unpark();
+    }
+}
+
+impl<'a> ::executor::Sleep for &'a Arc<ThreadNotify> {
+    type Wakeup = Arc<ThreadNotify>;
+
+    fn wakeup(&self) -> Self::Wakeup {
+        (*self).clone()
+    }
+
+    fn sleep(&mut self) {
+        self.park();
+    }
+
+    fn sleep_timeout(&mut self, duration: Duration) {
+        self.park_timeout(Some(duration));
+    }
+}
+
+impl ::executor::Wakeup for Arc<ThreadNotify> {
+    fn wakeup(&self) {
+        self.unpark();
+    }
+}
+
 impl fmt::Debug for ThreadNotify {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         fmt.debug_struct("ThreadNotify").finish()


Test patch:
diff --git a/tests/current_thread_executor.rs b/tests/current_thread_executor.rs
index fc9f2701..e96713e9 100644
--- a/tests/current_thread_executor.rs
+++ b/tests/current_thread_executor.rs
@@ -11,7 +11,7 @@ use std::rc::Rc;
 fn spawning_from_init_future() {
     let cnt = Rc::new(Cell::new(0));
 
-    CurrentThread::block_with_init(|_| {
+    CurrentThread::run(|_| {
         let cnt = cnt.clone();
 
         CurrentThread::execute(lazy(move || {
@@ -31,7 +31,7 @@ fn block_waits_for_non_daemon() {
 
     let cnt = Rc::new(Cell::new(0));
 
-    CurrentThread::block_with_init(|_| {
+    CurrentThread::run(|_| {
         let cnt = cnt.clone();
 
         let (tx, rx) = oneshot::channel();
@@ -62,7 +62,7 @@ fn spawn_many() {
 
     let cnt = Rc::new(Cell::new(0));
 
-    CurrentThread::block_with_init(|_| {
+    CurrentThread::run(|_| {
         for _ in 0..ITER {
             let cnt = cnt.clone();
             CurrentThread::execute(lazy(move || {
@@ -90,7 +90,7 @@ impl Future for Never {
 fn outstanding_daemon_tasks_are_dropped_on_return() {
     let mut rc = Rc::new(());
 
-    CurrentThread::block_with_init(|_| {
+    CurrentThread::run(|_| {
         CurrentThread::execute_daemon(Never(rc.clone()));
     });
 
@@ -102,7 +102,7 @@ fn outstanding_daemon_tasks_are_dropped_on_return() {
 fn outstanding_tasks_are_dropped_on_cancel() {
     let mut rc = Rc::new(());
 
-    CurrentThread::block_with_init(|_| {
+    CurrentThread::run(|_| {
         CurrentThread::execute(Never(rc.clone()));
         CurrentThread::cancel_all_executing();
     });
@@ -113,19 +113,19 @@ fn outstanding_tasks_are_dropped_on_cancel() {
 
 #[test]
 #[should_panic]
-fn nesting_block_with_init() {
-    CurrentThread::block_with_init(|_| {
-        CurrentThread::block_with_init(|_| {
+fn nesting_run() {
+    CurrentThread::run(|_| {
+        CurrentThread::run(|_| {
         });
     });
 }
 
 #[test]
 #[should_panic]
-fn block_with_init_in_future() {
-    CurrentThread::block_with_init(|_| {
+fn run_in_future() {
+    CurrentThread::run(|_| {
         CurrentThread::execute(lazy(|| {
-            CurrentThread::block_with_init(|_| {
+            CurrentThread::run(|_| {
             });
             Ok::<(), ()>(())
         }));
@@ -135,7 +135,7 @@ fn block_with_init_in_future() {
 #[test]
 #[should_panic]
 fn blocking_within_init() {
-    CurrentThread::block_with_init(|_| {
+    CurrentThread::run(|_| {
         let _ = blocking(empty::<(), ()>()).wait();
     });
 }
@@ -143,7 +143,7 @@ fn blocking_within_init() {
 #[test]
 #[should_panic]
 fn blocking_in_future() {
-    CurrentThread::block_with_init(|_| {
+    CurrentThread::run(|_| {
         CurrentThread::execute(lazy(|| {
             let _ = blocking(empty::<(), ()>()).wait();
             Ok::<(), ()>(())
@@ -184,7 +184,7 @@ fn tasks_are_scheduled_fairly() {
         }
     }
 
-    CurrentThread::block_with_init(|_| {
+    CurrentThread::run(|_| {
         CurrentThread::execute(Spin {
             state: state.clone(),
             idx: 0,