macros: add "local" runtime flavor

Author: Twey

tests-integration/tests/macros_main.rs

@@ -16,6 +16,13 @@ async fn spawning() -> usize {
     join.await.unwrap()
 }
 
+#[cfg(tokio_unstable)]
+#[tokio::main(flavor = "local")]
+async fn local_main() -> usize {
+    let join = tokio::task::spawn_local(async { 1 });
+    join.await.unwrap()
+}
+
 #[test]
 fn main_with_spawn() {
     assert_eq!(1, spawning());
@@ -24,5 +31,8 @@ fn main_with_spawn() {
 #[test]
 fn shell() {
     assert_eq!(1, basic_main());
-    assert_eq!(bool::default(), generic_fun::<bool>())
+    assert_eq!(bool::default(), generic_fun::<bool>());
+
+    #[cfg(tokio_unstable)]
+    assert_eq!(1, local_main());
 }

tokio-macros/src/entry.rs

@@ -10,13 +10,15 @@ type AttributeArgs = syn::punctuated::Punctuated<syn::Meta, syn::Token![,]>;
 enum RuntimeFlavor {
     CurrentThread,
     Threaded,
+    Local,
 }
 
 impl RuntimeFlavor {
     fn from_str(s: &str) -> Result<RuntimeFlavor, String> {
         match s {
             "current_thread" => Ok(RuntimeFlavor::CurrentThread),
             "multi_thread" => Ok(RuntimeFlavor::Threaded),
+            "local" => Ok(RuntimeFlavor::Local),
             "single_thread" => Err("The single threaded runtime flavor is called `current_thread`.".to_string()),
             "basic_scheduler" => Err("The `basic_scheduler` runtime flavor has been renamed to `current_thread`.".to_string()),
             "threaded_scheduler" => Err("The `threaded_scheduler` runtime flavor has been renamed to `multi_thread`.".to_string()),
@@ -178,14 +180,14 @@ impl Configuration {
 
         let flavor = self.flavor.unwrap_or(self.default_flavor);
         let worker_threads = match (flavor, self.worker_threads) {
-            (F::CurrentThread, Some((_, worker_threads_span))) => {
+            (F::CurrentThread | F::Local, Some((_, worker_threads_span))) => {
                 let msg = format!(
                     "The `worker_threads` option requires the `multi_thread` runtime flavor. Use `#[{}(flavor = \"multi_thread\")]`",
                     self.macro_name(),
                 );
                 return Err(syn::Error::new(worker_threads_span, msg));
             }
-            (F::CurrentThread, None) => None,
+            (F::CurrentThread | F::Local, None) => None,
             (F::Threaded, worker_threads) if self.rt_multi_thread_available => {
                 worker_threads.map(|(val, _span)| val)
             }
@@ -207,7 +209,7 @@ impl Configuration {
                 );
                 return Err(syn::Error::new(start_paused_span, msg));
             }
-            (F::CurrentThread, Some((start_paused, _))) => Some(start_paused),
+            (F::CurrentThread | F::Local, Some((start_paused, _))) => Some(start_paused),
             (_, None) => None,
         };
 
@@ -219,7 +221,7 @@ impl Configuration {
                 );
                 return Err(syn::Error::new(unhandled_panic_span, msg));
             }
-            (F::CurrentThread, Some((unhandled_panic, _))) => Some(unhandled_panic),
+            (F::CurrentThread | F::Local, Some((unhandled_panic, _))) => Some(unhandled_panic),
             (_, None) => None,
         };
 
@@ -408,13 +410,22 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
         .unwrap_or_else(|| Ident::new("tokio", last_stmt_start_span).into_token_stream());
 
     let mut rt = match config.flavor {
-        RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_start_span=>
-            #crate_path::runtime::Builder::new_current_thread()
-        },
+        RuntimeFlavor::CurrentThread | RuntimeFlavor::Local => {
+            quote_spanned! {last_stmt_start_span=>
+                #crate_path::runtime::Builder::new_current_thread()
+            }
+        }
         RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=>
             #crate_path::runtime::Builder::new_multi_thread()
         },
     };
+
+    let build = if let RuntimeFlavor::Local = config.flavor {
+        quote_spanned! {last_stmt_start_span=> build_local(Default::default())}
+    } else {
+        quote_spanned! {last_stmt_start_span=> build()}
+    };
+
     if let Some(v) = config.worker_threads {
         rt = quote_spanned! {last_stmt_start_span=> #rt.worker_threads(#v) };
     }
@@ -441,7 +452,7 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
         {
             return #rt
                 .enable_all()
-                .build()
+                .#build
                 .expect("Failed building the Runtime")
                 .block_on(#body_ident);
         }