macros: run runtime inside LocalSet when using macro (#4027)

tokio-rs · Sep 15, 2021 · 33f0a1f · 33f0a1f
1 parent 4c9b469
commit 33f0a1f
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 34 deletions.
diff --git a/tests-build/tests/fail/macros_type_mismatch.stderr b/tests-build/tests/fail/macros_type_mismatch.stderr
@@ -1,8 +1,8 @@
 error[E0308]: mismatched types
- --> $DIR/macros_type_mismatch.rs:5:5
+ --> $DIR/macros_type_mismatch.rs:5:7
   |
 5 |     Ok(())
-  |     ^^^^^^ expected `()`, found enum `Result`
+  |       ^^^^ expected `()`, found enum `Result`
   |
   = note: expected unit type `()`
                   found enum `Result<(), _>`
@@ -16,12 +16,12 @@ help: try adding a return type
   |                                             ^^^^^^^^^^^^^^^^
 
 error[E0308]: mismatched types
-  --> $DIR/macros_type_mismatch.rs:10:5
+  --> $DIR/macros_type_mismatch.rs:10:18
    |
 9  | async fn missing_return_type() {
    |                                - help: try adding a return type: `-> Result<(), _>`
 10 |     return Ok(());
-   |     ^^^^^^^^^^^^^^ expected `()`, found enum `Result`
+   |                  ^ expected `()`, found enum `Result`
    |
    = note: expected unit type `()`
                    found enum `Result<(), _>`

diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs
@@ -339,11 +339,10 @@ fn parse_knobs(
         {
             let body = async #body;
             #[allow(clippy::expect_used)]
-            #tail_return #rt
-                .enable_all()
-                .build()
-                .expect("Failed building the Runtime")
-                .block_on(body)#tail_semicolon
+            #tail_return tokio::task::LocalSet::new().block_on(
+              &#rt.enable_all().build().expect("Failed building the Runtime"),
+              body,
+            )#tail_semicolon
         }
     })
     .expect("Parsing failure");

diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs
@@ -27,6 +27,9 @@ use proc_macro::TokenStream;
 /// helps set up a `Runtime` without requiring the user to use
 /// [Runtime](../tokio/runtime/struct.Runtime.html) or
 /// [Builder](../tokio/runtime/struct.Builder.html) directly.
+/// 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.
 ///
 /// Note: This macro is designed to be simplistic and targets applications that
 /// do not require a complex setup. If the provided functionality is not
@@ -84,13 +87,14 @@ use proc_macro::TokenStream;
 ///
 /// ```rust
 /// fn main() {
-///     tokio::runtime::Builder::new_multi_thread()
+///     let ls = tokio::task::LocalSet::new();
+///     let rt = tokio::runtime::Builder::new_multi_thread()
 ///         .enable_all()
 ///         .build()
-///         .unwrap()
-///         .block_on(async {
-///             println!("Hello world");
-///         })
+///         .unwrap();
+///     ls.block_on(&rt, async {
+///         println!("Hello world");
+///     })
 /// }
 /// ```
 ///
@@ -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");
+///     })
 /// }
 /// ```
 ///
@@ -132,14 +137,15 @@ use proc_macro::TokenStream;
 ///
 /// ```rust
 /// fn main() {
-///     tokio::runtime::Builder::new_multi_thread()
+///     let ls = tokio::task::LocalSet::new();
+///     let rt = tokio::runtime::Builder::new_multi_thread()
 ///         .worker_threads(2)
 ///         .enable_all()
 ///         .build()
-///         .unwrap()
-///         .block_on(async {
-///             println!("Hello world");
-///         })
+///         .unwrap();
+///     ls.block_on(&rt, async {
+///         println!("Hello world");
+///     })
 /// }
 /// ```
 ///
@@ -156,14 +162,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 +211,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");
+///     })
 /// }
 /// ```
 ///

diff --git a/tokio/tests/task_local_set.rs b/tokio/tests/task_local_set.rs
@@ -16,6 +16,16 @@ 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")]
+async fn localset_implicit_multi_thread() {
+    task::spawn_local(async {}).await.unwrap();
+}
+
 #[tokio::test(flavor = "current_thread")]
 async fn local_basic_scheduler() {
     LocalSet::new()