Async/await polish (#1058)

A general refresh of Tokio's experimental async / await support.

6 files changed, 42 insertions, 47 deletions

diff --git a/tokio/Cargo.toml b/tokio/Cargo.toml
index 6f186d5b..e4825532 100644
--- a/tokio/Cargo.toml
+++ b/tokio/Cargo.toml
@@ -58,7 +58,8 @@ uds = ["tokio-uds"]
 # This feature comes with no promise of stability. Things will
 # break with each patch release. Use at your own risk.
 async-await-preview = [
-  "tokio-async-await/async-await-preview",
+  "tokio-futures/async-await-preview",
+  "tokio-macros/async-await-preview",
 ]
 
 [dependencies]
@@ -73,6 +74,7 @@ tokio-current-thread = { version = "0.1.6", optional = true }
 tokio-fs = { version = "0.1.6", optional = true }
 tokio-io = { version = "0.1.6", optional = true }
 tokio-executor = { version = "0.1.7", optional = true }
+tokio-macros = { version = "0.1.0", optional = true, path = "../tokio-macros" }
 tokio-reactor = { version = "0.1.1", optional = true }
 tokio-sync = { version = "0.1.5", optional = true }
 tokio-threadpool = { version = "0.1.13", optional = true }
@@ -85,7 +87,7 @@ tokio-trace-core = { version = "0.1", optional = true }
 mio = { version = "0.6.14", optional = true }
 
 # Needed for async/await preview support
-tokio-async-await = { version = "0.1.0", optional = true }
+tokio-futures = { version = "0.1.0", optional = true, path = "../tokio-futures" }
 
 [target.'cfg(unix)'.dependencies]
 tokio-uds = { version = "0.2.1", optional = true }
diff --git a/tokio/src/async_await.rs b/tokio/src/async_await.rs
index 605d9264..ed8b52d0 100644
--- a/tokio/src/async_await.rs
+++ b/tokio/src/async_await.rs
@@ -1,48 +1,17 @@
-use std::future::Future as StdFuture;
-use std::pin::Pin;
-use std::task::{Context, Poll};
-
-fn map_ok<T: StdFuture>(future: T) -> impl StdFuture<Output = Result<(), ()>> {
-    MapOk(future)
-}
-
-struct MapOk<T>(T);
-
-impl<T> MapOk<T> {
-    fn future<'a>(self: Pin<&'a mut Self>) -> Pin<&'a mut T> {
-        unsafe { Pin::map_unchecked_mut(self, |x| &mut x.0) }
-    }
-}
-
-impl<T: StdFuture> StdFuture for MapOk<T> {
-    type Output = Result<(), ()>;
-
-    fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
-        match self.future().poll(context) {
-            Poll::Ready(_) => Poll::Ready(Ok(())),
-            Poll::Pending => Poll::Pending,
-        }
-    }
-}
+use tokio_futures::compat;
 
 /// Like `tokio::run`, but takes an `async` block
 pub fn run_async<F>(future: F)
 where
-    F: StdFuture<Output = ()> + Send + 'static,
+    F: std::future::Future<Output = ()> + Send + 'static,
 {
-    use tokio_async_await::compat::backward;
-    let future = backward::Compat::new(map_ok(future));
-
-    ::run(future);
+    ::run(compat::infallible_into_01(future));
 }
 
 /// Like `tokio::spawn`, but takes an `async` block
 pub fn spawn_async<F>(future: F)
 where
-    F: StdFuture<Output = ()> + Send + 'static,
+    F: std::future::Future<Output = ()> + Send + 'static,
 {
-    use tokio_async_await::compat::backward;
-    let future = backward::Compat::new(map_ok(future));
-
-    ::spawn(future);
+    ::spawn(compat::infallible_into_01(future));
 }
diff --git a/tokio/src/lib.rs b/tokio/src/lib.rs
index 90b70321..872f23cd 100644
--- a/tokio/src/lib.rs
+++ b/tokio/src/lib.rs
@@ -1,9 +1,6 @@
 #![doc(html_root_url = "https://docs.rs/tokio/0.1.19")]
 #![deny(missing_docs, warnings, missing_debug_implementations)]
-#![cfg_attr(
-    feature = "async-await-preview",
-    feature(async_await, await_macro, futures_api,)
-)]
+#![cfg_attr(feature = "async-await-preview", feature(async_await, await_macro))]
 
 //! A runtime for writing reliable, asynchronous, and slim applications.
 //!
@@ -108,9 +105,6 @@ extern crate tokio_timer;
 #[cfg(feature = "udp")]
 extern crate tokio_udp;
 
-#[cfg(feature = "async-await-preview")]
-extern crate tokio_async_await;
-
 #[cfg(all(unix, feature = "uds"))]
 extern crate tokio_uds;
 
@@ -146,10 +140,19 @@ if_runtime! {
 // ===== Experimental async/await support =====
 
 #[cfg(feature = "async-await-preview")]
+extern crate tokio_futures;
+
+#[cfg(feature = "async-await-preview")]
+extern crate tokio_macros;
+
+#[cfg(feature = "async-await-preview")]
 mod async_await;
 
 #[cfg(feature = "async-await-preview")]
 pub use async_await::{run_async, spawn_async};
 
 #[cfg(feature = "async-await-preview")]
-pub use tokio_async_await::await;
+pub use tokio_futures::await;
+
+#[cfg(feature = "async-await-preview")]
+pub use tokio_macros::{main, test};
diff --git a/tokio/src/prelude.rs b/tokio/src/prelude.rs
index 17b2469d..b1b5f85b 100644
--- a/tokio/src/prelude.rs
+++ b/tokio/src/prelude.rs
@@ -21,7 +21,7 @@ pub use futures::{future, stream, task, Async, AsyncSink, Future, IntoFuture, Po
 
 #[cfg(feature = "async-await-preview")]
 #[doc(inline)]
-pub use tokio_async_await::{
+pub use tokio_futures::{
     io::{AsyncReadExt, AsyncWriteExt},
     sink::SinkExt,
     stream::StreamExt as StreamAsyncExt,
diff --git a/tokio/src/runtime/threadpool/async_await.rs b/tokio/src/runtime/threadpool/async_await.rs
new file mode 100644
index 00000000..12bdb1a1
--- /dev/null
+++ b/tokio/src/runtime/threadpool/async_await.rs
@@ -0,0 +1,18 @@
+use super::Runtime;
+use std::future::Future;
+
+impl Runtime {
+    /// Like `block_on`, but takes an `async` block
+    pub fn block_on_async<F>(&mut self, future: F) -> F::Output
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        use tokio_futures::compat;
+
+        match self.block_on(compat::infallible_into_01(future)) {
+            Ok(v) => v,
+            Err(_) => unreachable!(),
+        }
+    }
+}
diff --git a/tokio/src/runtime/threadpool/mod.rs b/tokio/src/runtime/threadpool/mod.rs
index 77d89498..b688a467 100644
--- a/tokio/src/runtime/threadpool/mod.rs
+++ b/tokio/src/runtime/threadpool/mod.rs
@@ -2,6 +2,9 @@ mod builder;
 mod shutdown;
 mod task_executor;
 
+#[cfg(feature = "async-await-preview")]
+mod async_await;
+
 pub use self::builder::Builder;
 pub use self::shutdown::Shutdown;
 pub use self::task_executor::TaskExecutor;