Simultaneous futures compat (#172)

This patch adds opt-in support for futures 0.2.

tokio-threadpool/Cargo.toml

@@ -19,8 +19,13 @@ crossbeam-deque = "0.3"
 num_cpus = "1.2"
 rand = "0.4"
 log = "0.3"
+futures2 = { version = "0.1", path = "../futures2", optional = true }
 
 [dev-dependencies]
 tokio-timer = "0.1"
 env_logger = "0.4"
 futures-cpupool = "0.1.7"
+
+[features]
+unstable-futures = ["futures2", "tokio-executor/unstable-futures"]
+default = []

tokio-threadpool/src/lib.rs

@@ -12,6 +12,9 @@ extern crate rand;
 #[macro_use]
 extern crate log;
 
+#[cfg(feature = "unstable-futures")]
+extern crate futures2;
+
 mod task;
 
 use tokio_executor::{Enter, SpawnError};
@@ -33,6 +36,14 @@ use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering::{AcqRel, Acquire, Release, Relaxed};
 use std::time::{Instant, Duration};
 
+#[derive(Debug)]
+struct ShutdownTask {
+    task1: AtomicTask,
+
+    #[cfg(feature = "unstable-futures")]
+    task2: futures2::task::AtomicWaker,
+}
+
 /// Work-stealing based thread pool for executing futures.
 ///
 /// If a `ThreadPool` instance is dropped without explicitly being shutdown,
@@ -160,7 +171,7 @@ struct Inner {
     workers: Box<[WorkerEntry]>,
 
     // Task notified when the worker shuts down
-    shutdown_task: AtomicTask,
+    shutdown_task: ShutdownTask,
 
     // Configuration
     config: Config,
@@ -180,6 +191,12 @@ struct Notifier {
     inner: Weak<Inner>,
 }
 
+#[cfg(feature = "unstable-futures")]
+struct Futures2Wake {
+    notifier: Arc<Notifier>,
+    id: usize,
+}
+
 /// ThreadPool state.
 ///
 /// The two least significant bits are the shutdown flags.  (0 for active, 1 for
@@ -532,7 +549,11 @@ impl Builder {
             num_workers: AtomicUsize::new(self.pool_size),
             next_thread_id: AtomicUsize::new(0),
             workers: workers.into_boxed_slice(),
-            shutdown_task: AtomicTask::new(),
+            shutdown_task: ShutdownTask {
+                task1: AtomicTask::new(),
+                #[cfg(feature = "unstable-futures")]
+                task2: futures2::task::AtomicWaker::new(),
+            },
             config: self.config.clone(),
         });
 
@@ -772,6 +793,11 @@ impl tokio_executor::Executor for Sender {
         let mut s = &*self;
         tokio_executor::Executor::spawn(&mut s, future)
     }
+
+    #[cfg(feature = "unstable-futures")]
+    fn spawn2(&mut self, f: Task2) -> Result<(), futures2::executor::SpawnError> {
+        futures2::executor::Executor::spawn(self, f)
+    }
 }
 
 impl<'a> tokio_executor::Executor for &'a Sender {
@@ -806,6 +832,11 @@ impl<'a> tokio_executor::Executor for &'a Sender {
 
         Ok(())
     }
+
+    #[cfg(feature = "unstable-futures")]
+    fn spawn2(&mut self, f: Task2) -> Result<(), futures2::executor::SpawnError> {
+        futures2::executor::Executor::spawn(self, f)
+    }
 }
 
 impl<T> future::Executor<T> for Sender
@@ -827,6 +858,48 @@ where T: Future<Item = (), Error = ()> + Send + 'static,
     }
 }
 
+#[cfg(feature = "unstable-futures")]
+type Task2 = Box<futures2::Future<Item = (), Error = futures2::Never> + Send>;
+
+#[cfg(feature = "unstable-futures")]
+impl futures2::executor::Executor for Sender {
+    fn spawn(&mut self, f: Task2) -> Result<(), futures2::executor::SpawnError> {
+        let mut s = &*self;
+        futures2::executor::Executor::spawn(&mut s, f)
+    }
+
+    fn status(&self) -> Result<(), futures2::executor::SpawnError> {
+        let s = &*self;
+        futures2::executor::Executor::status(&s)
+    }
+}
+
+#[cfg(feature = "unstable-futures")]
+impl<'a> futures2::executor::Executor for &'a Sender {
+    fn spawn(&mut self, f: Task2) -> Result<(), futures2::executor::SpawnError> {
+        self.prepare_for_spawn()
+            // TODO: get rid of this once the futures crate adds more error types
+            .map_err(|_| futures2::executor::SpawnError::shutdown())?;
+
+        // At this point, the pool has accepted the future, so schedule it for
+        // execution.
+
+        // Create a new task for the future
+        let task = Task::new2(f, |id| into_waker(Arc::new(Futures2Wake::new(id, &self.inner))));
+
+        self.inner.submit(task, &self.inner);
+
+        Ok(())
+    }
+
+    fn status(&self) -> Result<(), futures2::executor::SpawnError> {
+        tokio_executor::Executor::status(self)
+        // TODO: get rid of this once the futures crate adds more error types
+            .map_err(|_| futures2::executor::SpawnError::shutdown())
+    }
+}
+
+
 impl Clone for Sender {
     #[inline]
     fn clone(&self) -> Sender {
@@ -835,6 +908,21 @@ impl Clone for Sender {
     }
 }
 
+// ===== impl ShutdownTask =====
+
+impl ShutdownTask {
+    #[cfg(not(feature = "unstable-futures"))]
+    fn notify(&self) {
+        self.task1.notify();
+    }
+
+    #[cfg(feature = "unstable-futures")]
+    fn notify(&self) {
+        self.task1.notify();
+        self.task2.wake();
+    }
+}
+
 // ===== impl Shutdown =====
 
 impl Shutdown {
@@ -850,7 +938,7 @@ impl Future for Shutdown {
     fn poll(&mut self) -> Poll<(), ()> {
         trace!("Shutdown::poll");
 
-        self.inner().shutdown_task.register();
+        self.inner().shutdown_task.task1.register();
 
         if 0 != self.inner().num_workers.load(Acquire) {
             return Ok(Async::NotReady);
@@ -860,6 +948,24 @@ impl Future for Shutdown {
     }
 }
 
+#[cfg(feature = "unstable-futures")]
+impl futures2::Future for Shutdown {
+    type Item = ();
+    type Error = ();
+
+    fn poll(&mut self, cx: &mut futures2::task::Context) -> futures2::Poll<(), ()> {
+        trace!("Shutdown::poll");
+
+        self.inner().shutdown_task.task2.register(cx.waker());
+
+        if 0 != self.inner().num_workers.load(Acquire) {
+            return Ok(futures2::Async::Pending);
+        }
+
+        Ok(().into())
+    }
+}
+
 // ===== impl Inner =====
 
 impl Inner {
@@ -1346,6 +1452,7 @@ impl Worker {
         let notify = Arc::new(Notifier {
             inner: Arc::downgrade(&self.inner),
         });
+        let mut sender = Sender { inner: self.inner.clone() };
 
         let mut first = true;
         let mut spin_cnt = 0;
@@ -1358,14 +1465,14 @@ impl Worker {
             let consistent = self.drain_inbound();
 
             // Run the next available task
-            if self.try_run_task(&notify) {
+            if self.try_run_task(&notify, &mut sender) {
                 spin_cnt = 0;
                 // As long as there is work, keep looping.
                 continue;
             }
 
             // No work in this worker's queue, it is time to try stealing.
-            if self.try_steal_task(&notify) {
+            if self.try_steal_task(&notify, &mut sender) {
                 spin_cnt = 0;
                 continue;
             }
@@ -1448,13 +1555,13 @@ impl Worker {
     ///
     /// Returns `true` if work was found.
     #[inline]
-    fn try_run_task(&self, notify: &Arc<Notifier>) -> bool {
+    fn try_run_task(&self, notify: &Arc<Notifier>, sender: &mut Sender) -> bool {
         use deque::Steal::*;
 
         // Poll the internal queue for a task to run
         match self.entry().deque.steal() {
             Data(task) => {
-                self.run_task(task, notify);
+                self.run_task(task, notify, sender);
                 true
             }
             Empty => false,
@@ -1466,7 +1573,7 @@ impl Worker {
     ///
     /// Returns `true` if work was found
     #[inline]
-    fn try_steal_task(&self, notify: &Arc<Notifier>) -> bool {
+    fn try_steal_task(&self, notify: &Arc<Notifier>, sender: &mut Sender) -> bool {
         use deque::Steal::*;
 
         let len = self.inner.workers.len();
@@ -1480,7 +1587,7 @@ impl Worker {
                     Data(task) => {
                         trace!("stole task");
 
-                        self.run_task(task, notify);
+                        self.run_task(task, notify, sender);
 
                         trace!("try_steal_task -- signal_work; self={}; from={}",
                                self.idx, idx);
@@ -1507,10 +1614,10 @@ impl Worker {
         found_work
     }
 
-    fn run_task(&self, task: Task, notify: &Arc<Notifier>) {
+    fn run_task(&self, task: Task, notify: &Arc<Notifier>, sender: &mut Sender) {
         use task::Run::*;
 
-        match task.run(notify) {
+        match task.run(notify, sender) {
             Idle => {}
             Schedule => {
                 self.entry().push_internal(task);
@@ -2111,3 +2218,56 @@ impl fmt::Debug for Callback {
         write!(fmt, "Fn")
     }
 }
+
+// ===== impl Futures2Wake =====
+
+#[cfg(feature = "unstable-futures")]
+impl Futures2Wake {
+    fn new(id: usize, inner: &Arc<Inner>) -> Futures2Wake {
+        let notifier = Arc::new(Notifier {
+            inner: Arc::downgrade(inner),
+        });
+        Futures2Wake { id, notifier }
+    }
+}
+
+#[cfg(feature = "unstable-futures")]
+impl Drop for Futures2Wake {
+    fn drop(&mut self) {
+        self.notifier.drop_id(self.id)
+    }
+}
+
+#[cfg(feature = "unstable-futures")]
+struct ArcWrapped(PhantomData<Futures2Wake>);
+
+#[cfg(feature = "unstable-futures")]
+unsafe impl futures2::task::UnsafeWake for ArcWrapped {
+    unsafe fn clone_raw(&self) -> futures2::task::Waker {
+        let me: *const ArcWrapped = self;
+        let arc = (*(&me as *const *const ArcWrapped as *const Arc<Futures2Wake>)).clone();
+        arc.notifier.clone_id(arc.id);
+        into_waker(arc)
+    }
+
+    unsafe fn drop_raw(&self) {
+        let mut me: *const ArcWrapped = self;
+        let me = &mut me as *mut *const ArcWrapped as *mut Arc<Futures2Wake>;
+        (*me).notifier.drop_id((*me).id);
+        ::std::ptr::drop_in_place(me);
+    }
+
+    unsafe fn wake(&self) {
+        let me: *const ArcWrapped = self;
+        let me = &me as *const *const ArcWrapped as *const Arc<Futures2Wake>;
+        (*me).notifier.notify((*me).id)
+    }
+}
+
+#[cfg(feature = "unstable-futures")]
+fn into_waker(rc: Arc<Futures2Wake>) -> futures2::task::Waker {
+    unsafe {
+        let ptr = mem::transmute::<Arc<Futures2Wake>, *mut ArcWrapped>(rc);
+        futures2::task::Waker::new(ptr)
+    }
+}

tokio-threadpool/src/task.rs

@@ -1,6 +1,6 @@
-use Notifier;
+use {Notifier, Sender};
 
-use futures::{future, Future, Async};
+use futures::{self, future, Future, Async};
 use futures::executor::{self, Spawn};
 
 use std::{fmt, mem, panic, ptr};
@@ -9,6 +9,9 @@ use std::sync::Arc;
 use std::sync::atomic::{self, AtomicUsize, AtomicPtr};
 use std::sync::atomic::Ordering::{AcqRel, Acquire, Release, Relaxed};
 
+#[cfg(feature = "unstable-futures")]
+use futures2;
+
 pub(crate) struct Task {
     ptr: *mut Inner,
 }
@@ -34,6 +37,22 @@ pub(crate) enum Run {
     Complete,
 }
 
+type BoxFuture = Box<Future<Item = (), Error = ()> + Send + 'static>;
+
+#[cfg(feature = "unstable-futures")]
+type BoxFuture2 = Box<futures2::Future<Item = (), Error = futures2::Never> + Send>;
+
+enum TaskFuture {
+    Futures1(Spawn<BoxFuture>),
+
+    #[cfg(feature = "unstable-futures")]
+    Futures2 {
+        tls: futures2::task::LocalMap,
+        waker: futures2::task::Waker,
+        fut: BoxFuture2,
+    }
+}
+
 struct Inner {
     // Next pointer in the queue that submits tasks to a worker.
     next: AtomicPtr<Inner>,
@@ -47,7 +66,7 @@ struct Inner {
     // Store the future at the head of the struct
     //
     // The future is dropped immediately when it transitions to Complete
-    future: Option<Spawn<BoxFuture>>,
+    future: Option<TaskFuture>,
 }
 
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
@@ -64,23 +83,41 @@ enum State {
     Complete,
 }
 
-type BoxFuture = Box<Future<Item = (), Error = ()> + Send + 'static>;
-
 // ===== impl Task =====
 
 impl Task {
     /// Create a new task handle
     pub fn new(future: BoxFuture) -> Task {
+        let task_fut = TaskFuture::Futures1(executor::spawn(future));
         let inner = Box::new(Inner {
             next: AtomicPtr::new(ptr::null_mut()),
             state: AtomicUsize::new(State::new().into()),
             ref_count: AtomicUsize::new(1),
-            future: Some(executor::spawn(future)),
+            future: Some(task_fut),
         });
 
         Task { ptr: Box::into_raw(inner) }
     }
 
+    /// Create a new task handle for a futures 0.2 future
+    #[cfg(feature = "unstable-futures")]
+    pub fn new2<F>(fut: BoxFuture2, make_waker: F) -> Task
+        where F: FnOnce(usize) -> futures2::task::Waker
+    {
+        let mut inner = Box::new(Inner {
+            next: AtomicPtr::new(ptr::null_mut()),
+            state: AtomicUsize::new(State::new().into()),
+            ref_count: AtomicUsize::new(1),
+            future: None,
+        });
+
+        let waker = make_waker((&*inner) as *const _ as usize);
+        let tls = futures2::task::LocalMap::new();
+        inner.future = Some(TaskFuture::Futures2 { waker, tls, fut });
+
+        Task { ptr: Box::into_raw(inner) }
+    }
+
     /// Transmute a u64 to a Task
     pub unsafe fn from_notify_id(unpark_id: usize) -> Task {
         mem::transmute(unpark_id)
@@ -93,7 +130,7 @@ impl Task {
 
     /// Execute the task returning `Run::Schedule` if the task needs to be
     /// scheduled again.
-    pub fn run(&self, unpark: &Arc<Notifier>) -> Run {
+    pub fn run(&self, unpark: &Arc<Notifier>, exec: &mut Sender) -> Run {
         use self::State::*;
 
         // Transition task to running state. At this point, the task must be
@@ -118,7 +155,7 @@ impl Task {
         // `thread::panicking() -> true`. To do this, the future is dropped from
         // within the catch_unwind block.
         let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
-            struct Guard<'a>(&'a mut Option<Spawn<BoxFuture>>, bool);
+            struct Guard<'a>(&'a mut Option<TaskFuture>, bool);
 
             impl<'a> Drop for Guard<'a> {
                 fn drop(&mut self) {
@@ -132,7 +169,7 @@ impl Task {
             let mut g = Guard(fut, true);
 
             let ret = g.0.as_mut().unwrap()
-                .poll_future_notify(unpark, self.ptr as usize);
+                .poll(unpark, self.ptr as usize, exec);
 
 
             g.1 = false;
@@ -302,7 +339,7 @@ impl Inner {
             next: AtomicPtr::new(ptr::null_mut()),
             state: AtomicUsize::new(State::stub().into()),
             ref_count: AtomicUsize::new(0),
-            future: Some(executor::spawn(Box::new(future::empty()))),
+            future: Some(TaskFuture::Futures1(executor::spawn(Box::new(future::empty())))),
         }
     }
 
@@ -454,3 +491,23 @@ impl From<State> for usize {
         }
     }
 }
+
+// ===== impl TaskFuture =====
+
+impl TaskFuture {
+    #[allow(unused_variables)]
+    fn poll(&mut self, unpark: &Arc<Notifier>, id: usize, exec: &mut Sender) -> futures::Poll<(), ()> {
+        match *self {
+            TaskFuture::Futures1(ref mut fut) => fut.poll_future_notify(unpark, id),
+
+            #[cfg(feature = "unstable-futures")]
+            TaskFuture::Futures2 { ref mut fut, ref waker, ref mut tls } => {
+                let mut cx = futures2::task::Context::new(tls, waker, exec);
+                match fut.poll(&mut cx).unwrap() {
+                    futures2::Async::Pending => Ok(Async::NotReady),
+                    futures2::Async::Ready(x) => Ok(Async::Ready(x)),
+                }
+            }
+        }
+    }
+}

tokio-threadpool/tests/threadpool.rs

@@ -3,9 +3,20 @@ extern crate tokio_executor;
 extern crate futures;
 extern crate env_logger;
 
+#[cfg(feature = "unstable-futures")]
+extern crate futures2;
+
 use tokio_threadpool::*;
-use futures::{Poll, Sink, Stream, Async};
-use futures::future::{Future, lazy};
+
+#[cfg(not(feature = "unstable-futures"))]
+use futures::{Poll, Sink, Stream, Async, Future};
+#[cfg(not(feature = "unstable-futures"))]
+use futures::future::lazy;
+
+#[cfg(feature = "unstable-futures")]
+use futures2::prelude::*;
+#[cfg(feature = "unstable-futures")]
+use futures2::future::lazy;
 
 use std::cell::Cell;
 use std::sync::{mpsc, Arc};
@@ -15,6 +26,57 @@ use std::time::Duration;
 
 thread_local!(static FOO: Cell<u32> = Cell::new(0));
 
+#[cfg(not(feature = "unstable-futures"))]
+fn spawn_pool<F>(pool: &mut Sender, f: F)
+    where F: Future<Item = (), Error = ()> + Send + 'static
+{
+    pool.spawn(f).unwrap()
+}
+#[cfg(feature = "unstable-futures")]
+fn spawn_pool<F>(pool: &mut Sender, f: F)
+    where F: Future<Item = (), Error = ()> + Send + 'static
+{
+    futures2::executor::Executor::spawn(
+        pool,
+        Box::new(f.map_err(|_| panic!()))
+    ).unwrap()
+}
+
+#[cfg(not(feature = "unstable-futures"))]
+fn spawn_default<F>(f: F)
+    where F: Future<Item = (), Error = ()> + Send + 'static
+{
+    tokio_executor::spawn(f)
+}
+#[cfg(feature = "unstable-futures")]
+fn spawn_default<F>(f: F)
+    where F: Future<Item = (), Error = ()> + Send + 'static
+{
+    tokio_executor::spawn2(Box::new(f.map_err(|_| panic!())))
+}
+
+fn ignore_results<F: Future + Send + 'static>(f: F) -> Box<Future<Item = (), Error = ()> + Send> {
+    Box::new(f.map(|_| ()).map_err(|_| ()))
+}
+
+#[cfg(feature = "unstable-futures")]
+fn await_shutdown(shutdown: Shutdown) {
+    futures::Future::wait(shutdown).unwrap()
+}
+#[cfg(not(feature = "unstable-futures"))]
+fn await_shutdown(shutdown: Shutdown) {
+    shutdown.wait().unwrap()
+}
+
+#[cfg(not(feature = "unstable-futures"))]
+fn block_on<F: Future>(f: F) -> Result<F::Item, F::Error> {
+    f.wait()
+}
+#[cfg(feature = "unstable-futures")]
+fn block_on<F: Future>(f: F) -> Result<F::Item, F::Error> {
+    futures2::executor::block_on(f)
+}
+
 #[test]
 fn natural_shutdown_simple_futures() {
     let _ = ::env_logger::init();
@@ -33,29 +95,29 @@ fn natural_shutdown_simple_futures() {
                     NUM_DEC.fetch_add(1, Relaxed);
                 })
                 .build();
-            let tx = pool.sender().clone();
+            let mut tx = pool.sender().clone();
 
             let a = {
                 let (t, rx) = mpsc::channel();
-                tx.spawn(lazy(move || {
+                spawn_pool(&mut tx, lazy(move || {
                     // Makes sure this runs on a worker thread
                     FOO.with(|f| assert_eq!(f.get(), 0));
 
                     t.send("one").unwrap();
                     Ok(())
-                })).unwrap();
+                }));
                 rx
             };
 
             let b = {
                 let (t, rx) = mpsc::channel();
-                tx.spawn(lazy(move || {
+                spawn_pool(&mut tx, lazy(move || {
                     // Makes sure this runs on a worker thread
                     FOO.with(|f| assert_eq!(f.get(), 0));
 
                     t.send("two").unwrap();
                     Ok(())
-                })).unwrap();
+                }));
                 rx
             };
 
@@ -65,7 +127,7 @@ fn natural_shutdown_simple_futures() {
             assert_eq!("two", b.recv().unwrap());
 
             // Wait for the pool to shutdown
-            pool.shutdown().wait().unwrap();
+            await_shutdown(pool.shutdown());
 
             // Assert that at least one thread started
             let num_inc = NUM_INC.load(Relaxed);
@@ -89,6 +151,7 @@ fn force_shutdown_drops_futures() {
 
         struct Never(Arc<AtomicUsize>);
 
+        #[cfg(not(feature = "unstable-futures"))]
         impl Future for Never {
             type Item = ();
             type Error = ();
@@ -98,6 +161,16 @@ fn force_shutdown_drops_futures() {
             }
         }
 
+        #[cfg(feature = "unstable-futures")]
+        impl Future for Never {
+            type Item = ();
+            type Error = ();
+
+            fn poll(&mut self, _: &mut futures2::task::Context) -> Poll<(), ()> {
+                Ok(Async::Pending)
+            }
+        }
+
         impl Drop for Never {
             fn drop(&mut self) {
                 self.0.fetch_add(1, Relaxed);
@@ -116,10 +189,10 @@ fn force_shutdown_drops_futures() {
             .build();
         let mut tx = pool.sender().clone();
 
-        tx.spawn(Never(num_drop.clone())).unwrap();
+        spawn_pool(&mut tx, Never(num_drop.clone()));
 
         // Wait for the pool to shutdown
-        pool.shutdown_now().wait().unwrap();
+        await_shutdown(pool.shutdown_now());
 
         // Assert that only a single thread was spawned.
         let a = num_inc.load(Relaxed);
@@ -146,6 +219,7 @@ fn drop_threadpool_drops_futures() {
 
         struct Never(Arc<AtomicUsize>);
 
+        #[cfg(not(feature = "unstable-futures"))]
         impl Future for Never {
             type Item = ();
             type Error = ();
@@ -155,6 +229,16 @@ fn drop_threadpool_drops_futures() {
             }
         }
 
+        #[cfg(feature = "unstable-futures")]
+        impl Future for Never {
+            type Item = ();
+            type Error = ();
+
+            fn poll(&mut self, _: &mut futures2::task::Context) -> Poll<(), ()> {
+                Ok(Async::Pending)
+            }
+        }
+
         impl Drop for Never {
             fn drop(&mut self) {
                 self.0.fetch_add(1, Relaxed);
@@ -173,7 +257,7 @@ fn drop_threadpool_drops_futures() {
             .build();
         let mut tx = pool.sender().clone();
 
-        tx.spawn(Never(num_drop.clone())).unwrap();
+        spawn_pool(&mut tx, Never(num_drop.clone()));
 
         // Wait for the pool to shutdown
         drop(pool);
@@ -211,13 +295,13 @@ fn thread_shutdown_timeout() {
             let _ = t.lock().unwrap().send(());
         })
         .build();
-    let tx = pool.sender().clone();
+    let mut tx = pool.sender().clone();
 
     let t = complete_tx.clone();
-    tx.spawn(lazy(move || {
+    spawn_pool(&mut tx, lazy(move || {
         t.send(()).unwrap();
         Ok(())
-    })).unwrap();
+    }));
 
     // The future completes
     complete_rx.recv().unwrap();
@@ -226,14 +310,14 @@ fn thread_shutdown_timeout() {
     shutdown_rx.recv().unwrap();
 
     // Futures can still be run
-    tx.spawn(lazy(move || {
+    spawn_pool(&mut tx, lazy(move || {
         complete_tx.send(()).unwrap();
         Ok(())
-    })).unwrap();
+    }));
 
     complete_rx.recv().unwrap();
 
-    pool.shutdown().wait().unwrap();
+    await_shutdown(pool.shutdown());
 }
 
 #[test]
@@ -249,14 +333,14 @@ fn many_oneshot_futures() {
 
         for _ in 0..NUM {
             let cnt = cnt.clone();
-            tx.spawn(lazy(move || {
+            spawn_pool(&mut tx, lazy(move || {
                 cnt.fetch_add(1, Relaxed);
                 Ok(())
-            })).unwrap();
+            }));
         }
 
         // Wait for the pool to shutdown
-        pool.shutdown().wait().unwrap();
+        await_shutdown(pool.shutdown());
 
         let num = cnt.load(Relaxed);
         assert_eq!(num, NUM);
@@ -265,8 +349,12 @@ fn many_oneshot_futures() {
 
 #[test]
 fn many_multishot_futures() {
+    #[cfg(not(feature = "unstable-futures"))]
     use futures::sync::mpsc;
 
+    #[cfg(feature = "unstable-futures")]
+    use futures2::channel::mpsc;
+
     const CHAIN: usize = 200;
     const CYCLES: usize = 5;
     const TRACKS: usize = 50;
@@ -290,11 +378,11 @@ fn many_multishot_futures() {
                     .map_err(|e| panic!("{:?}", e));
 
                 // Forward all the messages
-                pool_tx.spawn(next_tx
+                spawn_pool(&mut pool_tx, next_tx
                     .send_all(rx)
                     .map(|_| ())
                     .map_err(|e| panic!("{:?}", e))
-                ).unwrap();
+                );
 
                 chain_rx = next_rx;
             }
@@ -304,7 +392,7 @@ fn many_multishot_futures() {
             let cycle_tx = start_tx.clone();
             let mut rem = CYCLES;
 
-            pool_tx.spawn(chain_rx.take(CYCLES as u64).for_each(move |msg| {
+            let task = chain_rx.take(CYCLES as u64).for_each(move |msg| {
                 rem -= 1;
                 let send = if rem == 0 {
                     final_tx.clone().send(msg)
@@ -316,83 +404,109 @@ fn many_multishot_futures() {
                     res.unwrap();
                     Ok(())
                 })
-            })).unwrap();
+            });
+            spawn_pool(&mut pool_tx, ignore_results(task));
 
             start_txs.push(start_tx);
             final_rxs.push(final_rx);
         }
 
         for start_tx in start_txs {
-            start_tx.send("ping").wait().unwrap();
+            block_on(start_tx.send("ping")).unwrap();
         }
 
         for final_rx in final_rxs {
-            final_rx.wait().next().unwrap().unwrap();
+            block_on(final_rx.into_future()).unwrap();
         }
 
         // Shutdown the pool
-        pool.shutdown().wait().unwrap();
+        await_shutdown(pool.shutdown());
     }
 }
 
 #[test]
 fn global_executor_is_configured() {
     let pool = ThreadPool::new();
-    let tx = pool.sender().clone();
+    let mut tx = pool.sender().clone();
 
     let (signal_tx, signal_rx) = mpsc::channel();
 
-    tx.spawn(lazy(move || {
-        tokio_executor::spawn(lazy(move || {
+    spawn_pool(&mut tx, lazy(move || {
+        spawn_default(lazy(move || {
             signal_tx.send(()).unwrap();
             Ok(())
         }));
 
         Ok(())
-    })).unwrap();
+    }));
 
     signal_rx.recv().unwrap();
 
-    pool.shutdown().wait().unwrap();
+    await_shutdown(pool.shutdown());
 }
 
 #[test]
 fn new_threadpool_is_idle() {
     let pool = ThreadPool::new();
-    pool.shutdown_on_idle().wait().unwrap();
+    await_shutdown(pool.shutdown_on_idle());
 }
 
 #[test]
 fn busy_threadpool_is_not_idle() {
+    #[cfg(not(feature = "unstable-futures"))]
     use futures::sync::oneshot;
 
+    #[cfg(feature = "unstable-futures")]
+    use futures2::channel::oneshot;
+
     let pool = ThreadPool::new();
-    let tx = pool.sender().clone();
+    let mut tx = pool.sender().clone();
 
     let (term_tx, term_rx) = oneshot::channel();
 
-    tx.spawn(term_rx.then(|_| {
+    spawn_pool(&mut tx, term_rx.then(|_| {
         Ok(())
-    })).unwrap();
+    }));
 
     let mut idle = pool.shutdown_on_idle();
 
-    futures::lazy(|| {
-        assert!(idle.poll().unwrap().is_not_ready());
-        Ok::<_, ()>(())
-    }).wait().unwrap();
+    struct IdleFut<'a>(&'a mut Shutdown);
+
+    #[cfg(not(feature = "unstable-futures"))]
+    impl<'a> Future for IdleFut<'a> {
+        type Item = ();
+        type Error = ();
+        fn poll(&mut self) -> Poll<(), ()> {
+            assert!(self.0.poll().unwrap().is_not_ready());
+            Ok(Async::Ready(()))
+        }
+    }
+
+    #[cfg(feature = "unstable-futures")]
+    impl<'a> Future for IdleFut<'a> {
+        type Item = ();
+        type Error = ();
+        fn poll(&mut self, cx: &mut futures2::task::Context) -> Poll<(), ()> {
+            assert!(self.0.poll(cx).unwrap().is_pending());
+            Ok(Async::Ready(()))
+        }
+    }
+
+    block_on(IdleFut(&mut idle)).unwrap();
 
     term_tx.send(()).unwrap();
 
-    idle.wait().unwrap();
+    await_shutdown(idle);
 }
 
 #[test]
 fn panic_in_task() {
     let pool = ThreadPool::new();
+    let mut tx = pool.sender().clone();
 
     struct Boom;
 
+    #[cfg(not(feature = "unstable-futures"))]
     impl Future for Boom {
         type Item = ();
         type Error = ();
@@ -402,13 +516,23 @@ fn panic_in_task() {
         }
     }
 
+    #[cfg(feature = "unstable-futures")]
+    impl Future for Boom {
+        type Item = ();
+        type Error = ();
+
+        fn poll(&mut self, _cx: &mut futures2::task::Context) -> Poll<(), ()> {
+            panic!();
+        }
+    }
+
     impl Drop for Boom {
         fn drop(&mut self) {
             assert!(::std::thread::panicking());
         }
     }
 
-    pool.spawn(Boom);
+    spawn_pool(&mut tx, Boom);
 
-    pool.shutdown_on_idle().wait().unwrap();
+    await_shutdown(pool.shutdown_on_idle());
 }