Simultaneous futures compat (#172)

This patch adds opt-in support for futures 0.2.

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());
 }