From 18cef1901fc805e7ef3443538c0931836891eab7 Mon Sep 17 00:00:00 2001
From: Sean McArthur <sean@seanmonstar.com>
Date: Tue, 24 Sep 2019 11:39:36 -0700
Subject: [PATCH] tokio: add `rt-current-thread` optional feature

- Adds a minimum `rt-current-thread` optional feature that exports
  `tokio::runtime::current_thread`.
- Adds a `macros` optional feature to enable the `#[tokio::main]` and
  `#[tokio::test]` attributes.
- Adjusts `#[tokio::main]` macro to select a runtime "automatically" if
  a specific strategy isn't specified. Allows using the macro with only
  the rt-current-thread feature.
---
 tokio-macros/src/lib.rs  | 10 +++++++++-
 tokio/Cargo.toml         |  8 +++++++-
 tokio/src/lib.rs         |  7 ++++++-
 tokio/src/runtime/mod.rs | 14 ++++++++++++++
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs
index 9a33c5fd8..d8eedb02a 100644
--- a/tokio-macros/src/lib.rs
+++ b/tokio-macros/src/lib.rs
@@ -49,6 +49,7 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
     enum RuntimeType {
         Single,
         Multi,
+        Auto,
     }
 
     let input = syn::parse_macro_input!(item as syn::ItemFn);
@@ -71,7 +72,7 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
             .into();
     }
 
-    let mut runtime = RuntimeType::Multi;
+    let mut runtime = RuntimeType::Auto;
 
     for arg in args {
         if let syn::NestedMeta::Meta(syn::Meta::Path(path)) = arg {
@@ -106,6 +107,13 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
                 rt.block_on(async { #body })
             }
         },
+        RuntimeType::Auto => quote! {
+            #(#attrs)*
+            fn #name() #ret {
+                let mut rt = tokio::runtime::__main::Runtime::new().unwrap();
+                rt.block_on(async { #body })
+            }
+        },
     };
 
     result.into()
diff --git a/tokio/Cargo.toml b/tokio/Cargo.toml
index 3728ee362..7b9a42810 100644
--- a/tokio/Cargo.toml
+++ b/tokio/Cargo.toml
@@ -37,15 +37,21 @@ default = [
 codec = ["io", "tokio-codec", "bytes"]
 fs = ["tokio-fs"]
 io = ["tokio-io"]
+macros = ["tokio-macros"]
 net = ["tcp", "udp", "uds"]
+rt-current-thread = [
+  "timer",
+  "tokio-net",
+  "tokio-executor/current-thread",
+]
 rt-full = [
+  "macros",
   "num_cpus",
   "net",
   "sync",
   "timer",
   "tokio-executor/current-thread",
   "tokio-executor/threadpool",
-  "tokio-macros",
   "tracing-core",
 ]
 signal = ["tokio-net/signal"]
diff --git a/tokio/src/lib.rs b/tokio/src/lib.rs
index 30375c0ad..14a7a3ab6 100644
--- a/tokio/src/lib.rs
+++ b/tokio/src/lib.rs
@@ -72,7 +72,10 @@
 
 macro_rules! if_runtime {
     ($($i:item)*) => ($(
-        #[cfg(any(feature = "rt-full"))]
+        #[cfg(any(
+            feature = "rt-full",
+            feature = "rt-current-thread",
+        ))]
         $i
     )*)
 }
@@ -103,8 +106,10 @@ if_runtime! {
     pub use crate::executor::spawn;
 
     #[cfg(not(test))] // Work around for rust-lang/rust#62127
+    #[cfg(feature = "macros")]
     #[doc(inline)]
     pub use tokio_macros::main;
+    #[cfg(feature = "macros")]
     #[doc(inline)]
     pub use tokio_macros::test;
 }
diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs
index ee33d08ca..aa183cbd3 100644
--- a/tokio/src/runtime/mod.rs
+++ b/tokio/src/runtime/mod.rs
@@ -135,10 +135,24 @@
 //! [`tokio::main`]: ../../tokio_macros/attr.main.html
 
 pub mod current_thread;
+#[cfg(feature = "rt-full")]
 mod threadpool;
 
+#[cfg(feature = "rt-full")]
 pub use self::threadpool::{
     Builder,
     Runtime,
     TaskExecutor,
 };
+
+// Internal export, don't use.
+// This exists to support "auto" runtime selection when using the
+// #[tokio::main] attribute.
+#[doc(hidden)]
+pub mod __main {
+    #[cfg(feature = "rt-full")]
+    pub use super::Runtime;
+
+    #[cfg(not(feature = "rt-full"))]
+    pub use super::current_thread::Runtime;
+}