macros: custom package name for #[tokio::main] and #[tokio::test]

kezhuw · kezhuw · commit fb47cccdeac3 · 2022-04-15T12:28:18.000+08:00
This also enables `#[crate::test(package = "crate")]` in unit tests. Sees rust-lang/cargo#5653. Fixes tokio-rs#2312.
diff --git a/tests-build/tests/fail/macros_invalid_input.rs b/tests-build/tests/fail/macros_invalid_input.rs
@@ -33,6 +33,15 @@ async fn test_worker_threads_not_int() {}
 #[tokio::test(flavor = "current_thread", worker_threads = 4)]
 async fn test_worker_threads_and_current_thread() {}
 
+#[tokio::test(package = 456)]
+async fn test_package_not_ident_int() {}
+
+#[tokio::test(package = "456")]
+async fn test_package_not_ident_invalid() {}
+
+#[tokio::test(package = "abc::edf")]
+async fn test_package_not_ident_path() {}
+
 #[tokio::test]
 #[test]
 async fn test_has_second_test_attr() {}
diff --git a/tests-build/tests/fail/macros_invalid_input.stderr b/tests-build/tests/fail/macros_invalid_input.stderr
@@ -4,7 +4,7 @@ error: the `async` keyword is missing from the function declaration
 4 | fn main_is_not_async() {}
   | ^^
 
-error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`
+error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `package`
  --> $DIR/macros_invalid_input.rs:6:15
   |
 6 | #[tokio::main(foo)]
@@ -22,13 +22,13 @@ error: the `async` keyword is missing from the function declaration
 13 | fn test_is_not_async() {}
    | ^^
 
-error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`
+error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `package`
   --> $DIR/macros_invalid_input.rs:15:15
    |
 15 | #[tokio::test(foo)]
    |               ^^^
 
-error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`
+error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `package`
   --> $DIR/macros_invalid_input.rs:18:15
    |
 18 | #[tokio::test(foo = 123)]
@@ -64,16 +64,34 @@ error: The `worker_threads` option requires the `multi_thread` runtime flavor. U
 33 | #[tokio::test(flavor = "current_thread", worker_threads = 4)]
    |                                                           ^
 
+error: Failed to parse value of `package` as ident.
+  --> $DIR/macros_invalid_input.rs:36:25
+   |
+36 | #[tokio::test(package = 456)]
+   |                         ^^^
+
+error: Failed to parse value of `package` as ident: "456"
+  --> $DIR/macros_invalid_input.rs:39:25
+   |
+39 | #[tokio::test(package = "456")]
+   |                         ^^^^^
+
+error: Failed to parse value of `package` as ident: "abc::edf"
+  --> $DIR/macros_invalid_input.rs:42:25
+   |
+42 | #[tokio::test(package = "abc::edf")]
+   |                         ^^^^^^^^^^
+
 error: second test attribute is supplied
-  --> $DIR/macros_invalid_input.rs:37:1
+  --> $DIR/macros_invalid_input.rs:46:1
    |
-37 | #[test]
+46 | #[test]
    | ^^^^^^^
 
 error: duplicated attribute
-  --> $DIR/macros_invalid_input.rs:37:1
+  --> $DIR/macros_invalid_input.rs:46:1
    |
-37 | #[test]
+46 | #[test]
    | ^^^^^^^
    |
    = note: `-D duplicate-macro-attributes` implied by `-D warnings`
diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs
@@ -1,5 +1,5 @@
 use proc_macro::TokenStream;
-use proc_macro2::Span;
+use proc_macro2::{Ident, Span};
 use quote::{quote, quote_spanned, ToTokens};
 use syn::parse::Parser;
 
@@ -29,13 +29,15 @@ struct FinalConfig {
     flavor: RuntimeFlavor,
     worker_threads: Option<usize>,
     start_paused: Option<bool>,
+    package_name: Option<String>,
 }
 
 /// Config used in case of the attribute not being able to build a valid config
 const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
     flavor: RuntimeFlavor::CurrentThread,
     worker_threads: None,
     start_paused: None,
+    package_name: None,
 };
 
 struct Configuration {
@@ -45,6 +47,7 @@ struct Configuration {
     worker_threads: Option<(usize, Span)>,
     start_paused: Option<(bool, Span)>,
     is_test: bool,
+    package_name: Option<String>,
 }
 
 impl Configuration {
@@ -59,6 +62,7 @@ impl Configuration {
             worker_threads: None,
             start_paused: None,
             is_test,
+            package_name: None,
         }
     }
 
@@ -104,6 +108,15 @@ impl Configuration {
         Ok(())
     }
 
+    fn set_package_name(&mut self, name: syn::Lit, span: Span) -> Result<(), syn::Error> {
+        if self.package_name.is_some() {
+            return Err(syn::Error::new(span, "`package` set multiple times."));
+        }
+        let name_ident = parse_ident(name, span, "package")?;
+        self.package_name = Some(name_ident.to_string());
+        Ok(())
+    }
+
     fn macro_name(&self) -> &'static str {
         if self.is_test {
             "tokio::test"
@@ -151,6 +164,7 @@ impl Configuration {
         };
 
         Ok(FinalConfig {
+            package_name: self.package_name.clone(),
             flavor,
             worker_threads,
             start_paused,
@@ -185,6 +199,27 @@ fn parse_string(int: syn::Lit, span: Span, field: &str) -> Result<String, syn::E
     }
 }
 
+fn parse_ident(lit: syn::Lit, span: Span, field: &str) -> Result<Ident, syn::Error> {
+    match lit {
+        syn::Lit::Str(s) => {
+            let err = syn::Error::new(
+                span,
+                format!(
+                    "Failed to parse value of `{}` as ident: \"{}\"",
+                    field,
+                    s.value()
+                ),
+            );
+            let path = s.parse::<syn::Path>().map_err(|_| err.clone())?;
+            path.get_ident().cloned().ok_or(err)
+        }
+        _ => Err(syn::Error::new(
+            span,
+            format!("Failed to parse value of `{}` as ident.", field),
+        )),
+    }
+}
+
 fn parse_bool(bool: syn::Lit, span: Span, field: &str) -> Result<bool, syn::Error> {
     match bool {
         syn::Lit::Bool(b) => Ok(b.value),
@@ -243,9 +278,15 @@ fn build_config(
                         let msg = "Attribute `core_threads` is renamed to `worker_threads`";
                         return Err(syn::Error::new_spanned(namevalue, msg));
                     }
+                    "package" => {
+                        config.set_package_name(
+                            namevalue.lit.clone(),
+                            syn::spanned::Spanned::span(&namevalue.lit),
+                        )?;
+                    }
                     name => {
                         let msg = format!(
-                            "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`",
+                            "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `package`",
                             name,
                         );
                         return Err(syn::Error::new_spanned(namevalue, msg));
@@ -275,7 +316,7 @@ fn build_config(
                         format!("The `{}` attribute requires an argument.", name)
                     }
                     name => {
-                        format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`", name)
+                        format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `package`", name)
                     }
                 };
                 return Err(syn::Error::new_spanned(path, msg));
@@ -313,12 +354,17 @@ fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> To
         (start, end)
     };
 
+    let package_name = config
+        .package_name
+        .map(|name| Ident::new(&name, last_stmt_start_span))
+        .unwrap_or_else(|| Ident::new("tokio", last_stmt_start_span));
+
     let mut rt = match config.flavor {
         RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_start_span=>
-            tokio::runtime::Builder::new_current_thread()
+            #package_name::runtime::Builder::new_current_thread()
         },
         RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=>
-            tokio::runtime::Builder::new_multi_thread()
+            #package_name::runtime::Builder::new_multi_thread()
         },
     };
     if let Some(v) = config.worker_threads {
diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs
@@ -168,12 +168,32 @@ use proc_macro::TokenStream;
 ///
 /// Note that `start_paused` requires the `test-util` feature to be enabled.
 ///
-/// ### NOTE:
+/// ### Rename package
 ///
-/// If you rename the Tokio crate in your dependencies this macro will not work.
-/// If you must rename the current version of Tokio because you're also using an
-/// older version of Tokio, you _must_ make the current version of Tokio
-/// available as `tokio` in the module where this macro is expanded.
+/// ```rust
+/// use tokio as tokio1;
+///
+/// #[tokio1::main(package = "tokio1")]
+/// async fn main() {
+///     println!("Hello world");
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```rust
+/// use tokio as tokio1;
+///
+/// fn main() {
+///     tokio1::runtime::Builder::new_multi_thread()
+///         .enable_all()
+///         .build()
+///         .unwrap()
+///         .block_on(async {
+///             println!("Hello world");
+///         })
+/// }
+/// ```
 #[proc_macro_attribute]
 #[cfg(not(test))] // Work around for rust-lang/rust#62127
 pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
@@ -213,12 +233,32 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
 /// }
 /// ```
 ///
-/// ### NOTE:
+/// ### Rename package
+///
+/// ```rust
+/// use tokio as tokio1;
+///
+/// #[tokio1::main(package = "tokio1")]
+/// async fn main() {
+///     println!("Hello world");
+/// }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```rust
+/// use tokio as tokio1;
 ///
-/// If you rename the Tokio crate in your dependencies this macro will not work.
-/// If you must rename the current version of Tokio because you're also using an
-/// older version of Tokio, you _must_ make the current version of Tokio
-/// available as `tokio` in the module where this macro is expanded.
+/// fn main() {
+///     tokio1::runtime::Builder::new_multi_thread()
+///         .enable_all()
+///         .build()
+///         .unwrap()
+///         .block_on(async {
+///             println!("Hello world");
+///         })
+/// }
+/// ```
 #[proc_macro_attribute]
 #[cfg(not(test))] // Work around for rust-lang/rust#62127
 pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream {
@@ -260,12 +300,16 @@ pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream {
 ///
 /// Note that `start_paused` requires the `test-util` feature to be enabled.
 ///
-/// ### NOTE:
+/// ### Rename package
+///
+/// ```rust
+/// use tokio as tokio1;
 ///
-/// If you rename the Tokio crate in your dependencies this macro will not work.
-/// If you must rename the current version of Tokio because you're also using an
-/// older version of Tokio, you _must_ make the current version of Tokio
-/// available as `tokio` in the module where this macro is expanded.
+/// #[tokio1::test(package = "tokio1")]
+/// async fn my_test() {
+///     println!("Hello world");
+/// }
+/// ```
 #[proc_macro_attribute]
 pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
     entry::test(args, item, true)
@@ -281,13 +325,6 @@ pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
 ///     assert!(true);
 /// }
 /// ```
-///
-/// ### NOTE:
-///
-/// If you rename the Tokio crate in your dependencies this macro will not work.
-/// If you must rename the current version of Tokio because you're also using an
-/// older version of Tokio, you _must_ make the current version of Tokio
-/// available as `tokio` in the module where this macro is expanded.
 #[proc_macro_attribute]
 pub fn test_rt(args: TokenStream, item: TokenStream) -> TokenStream {
     entry::test(args, item, false)
diff --git a/tokio/tests/macros_rename_test.rs b/tokio/tests/macros_rename_test.rs
@@ -0,0 +1,25 @@
+#![cfg(feature = "full")]
+
+use std as tokio;
+
+use ::tokio as tokio1;
+
+async fn compute() -> usize {
+    let join = tokio1::spawn(async { tokio::mem::size_of::<u8>() });
+    join.await.unwrap()
+}
+
+#[tokio1::main(package = "tokio1")]
+async fn compute_main() -> usize {
+    compute().await
+}
+
+#[test]
+fn package_rename_main() {
+    assert_eq!(1, compute_main());
+}
+
+#[tokio1::test(package = "tokio1")]
+async fn package_rename_test() {
+    assert_eq!(1, compute().await);
+}
