macros: run new_current_thread() inside LocalSet

tokio-macros/src/entry.rs

@@ -335,18 +335,33 @@ fn parse_knobs(
         ),
         _ => (quote! {}, quote! {}),
     };
-    input.block = syn::parse2(quote_spanned! {last_stmt_end_span=>
-        {
-            let body = async #body;
-            #[allow(clippy::expect_used)]
-            #tail_return #rt
-                .enable_all()
-                .build()
-                .expect("Failed building the Runtime")
-                .block_on(body)#tail_semicolon
-        }
-    })
-    .expect("Parsing failure");
+    let span = match config.flavor {
+        RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_end_span=>
+            {
+                let body = async #body;
+                #[allow(clippy::expect_used)]
+                #tail_return tokio::task::LocalSet::new().block_on(
+                    &#rt
+                        .enable_all()
+                        .build()
+                        .expect("Failed building the Runtime"),
+                    body,
+                )#tail_semicolon
+            }
+        },
+        RuntimeFlavor::Threaded => quote_spanned! {last_stmt_end_span=>
+            {
+                let body = async #body;
+                #[allow(clippy::expect_used)]
+                #tail_return #rt
+                    .enable_all()
+                    .build()
+                    .expect("Failed building the Runtime")
+                    .block_on(body)#tail_semicolon
+            }
+        },
+    };
+    input.block = syn::parse2(span).expect("Parsing failure");
     input.block.brace_token = brace_token;
 
     let result = quote! {

tokio-macros/src/lib.rs

@@ -98,6 +98,10 @@ use proc_macro::TokenStream;
 ///
 /// The basic scheduler is single-threaded.
 ///
+/// The function executes in the context of a
+/// [LocalSet](../tokio/task/struct.LocalSet.html), allowing calls to
+/// [spawn_local](../tokio/task/fn.spawn_local.html) without further setup.
+///
 /// ```rust
 /// #[tokio::main(flavor = "current_thread")]
 /// async fn main() {
@@ -109,13 +113,14 @@ use proc_macro::TokenStream;
 ///
 /// ```rust
 /// fn main() {
-///     tokio::runtime::Builder::new_current_thread()
+///     let ls = tokio::task::LocalSet::new();
+///     let rt = tokio::runtime::Builder::new_current_thread()
 ///         .enable_all()
 ///         .build()
-///         .unwrap()
-///         .block_on(async {
-///             println!("Hello world");
-///         })
+///         .unwrap();
+///     ls.block_on(&rt, async {
+///         println!("Hello world");
+///     })
 /// }
 /// ```
 ///
@@ -156,14 +161,15 @@ use proc_macro::TokenStream;
 ///
 /// ```rust
 /// fn main() {
-///     tokio::runtime::Builder::new_current_thread()
+///     let ls = tokio::task::LocalSet::new();
+///     let rt = tokio::runtime::Builder::new_current_thread()
 ///         .enable_all()
 ///         .start_paused(true)
 ///         .build()
-///         .unwrap()
-///         .block_on(async {
-///             println!("Hello world");
-///         })
+///         .unwrap();
+///     ls.block_on(&rt, async {
+///         println!("Hello world");
+///     })
 /// }
 /// ```
 ///
@@ -204,13 +210,14 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
 ///
 /// ```rust
 /// fn main() {
-///     tokio::runtime::Builder::new_current_thread()
+///     let ls = tokio::task::LocalSet::new();
+///     let rt = tokio::runtime::Builder::new_current_thread()
 ///         .enable_all()
 ///         .build()
-///         .unwrap()
-///         .block_on(async {
-///             println!("Hello world");
-///         })
+///         .unwrap();
+///     ls.block_on(&rt, async {
+///         println!("Hello world");
+///     })
 /// }
 /// ```
 ///

tokio/tests/task_local_set.rs

@@ -16,6 +16,17 @@ use std::sync::atomic::Ordering::{self, SeqCst};
 use std::sync::atomic::{AtomicBool, AtomicUsize};
 use std::time::Duration;
 
+#[tokio::test(flavor = "current_thread")]
+async fn localset_implicit_current_thread() {
+    task::spawn_local(async {}).await.unwrap();
+}
+
+#[tokio::test(flavor = "multi_thread")]
+#[should_panic]
+async fn localset_implicit_multi_thread() {
+    task::spawn_local(async {}).await.ok();
+}
+
 #[tokio::test(flavor = "current_thread")]
 async fn local_basic_scheduler() {
     LocalSet::new()