From d6f1654490f358ae81c8c115de4be91dfdf00619 Mon Sep 17 00:00:00 2001
From: ibraheemdev <ibrah1440@gmail.com>
Date: Tue, 8 Jun 2021 09:33:15 -0400
Subject: [PATCH] allow custom system path in actix_rt::main macro

---
 actix-macros/src/lib.rs                       | 42 +++++++++++++++++--
 actix-macros/tests/trybuild.rs                |  1 +
 .../tests/trybuild/main-04-system-path.rs     |  8 ++++
 3 files changed, 48 insertions(+), 3 deletions(-)
 create mode 100644 actix-macros/tests/trybuild/main-04-system-path.rs

diff --git a/actix-macros/src/lib.rs b/actix-macros/src/lib.rs
index 0d634a59..897d3339 100644
--- a/actix-macros/src/lib.rs
+++ b/actix-macros/src/lib.rs
@@ -27,8 +27,10 @@ use quote::quote;
 #[allow(clippy::needless_doctest_main)]
 #[proc_macro_attribute]
 #[cfg(not(test))] // Work around for rust-lang/rust#62127
-pub fn main(_: TokenStream, item: TokenStream) -> TokenStream {
+pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
     let mut input = syn::parse_macro_input!(item as syn::ItemFn);
+    let args = syn::parse_macro_input!(args as syn::AttributeArgs);
+
     let attrs = &input.attrs;
     let vis = &input.vis;
     let sig = &mut input.sig;
@@ -43,13 +45,47 @@ pub fn main(_: TokenStream, item: TokenStream) -> TokenStream {
         .into();
     }
 
+    let mut system = syn::parse_str::<syn::Path>("::actix_rt::System").unwrap();
+
+    for arg in &args {
+        match arg {
+            syn::NestedMeta::Meta(syn::Meta::NameValue(syn::MetaNameValue {
+                lit: syn::Lit::Str(lit),
+                path,
+                ..
+            })) => match path
+                .get_ident()
+                .map(|i| i.to_string().to_lowercase())
+                .as_deref()
+            {
+                Some("system") => match lit.parse() {
+                    Ok(path) => system = path,
+                    Err(_) => {
+                        return syn::Error::new_spanned(lit, "Expected path")
+                            .to_compile_error()
+                            .into();
+                    }
+                },
+                _ => {
+                    return syn::Error::new_spanned(arg, "Unkown attribute specified")
+                        .to_compile_error()
+                        .into();
+                }
+            },
+            _ => {
+                return syn::Error::new_spanned(arg, "Unknown attribute specified")
+                    .to_compile_error()
+                    .into();
+            }
+        }
+    }
+
     sig.asyncness = None;
 
     (quote! {
         #(#attrs)*
         #vis #sig {
-            actix_rt::System::new()
-                .block_on(async move { #body })
+            <#system>::new().block_on(async move { #body })
         }
     })
     .into()
diff --git a/actix-macros/tests/trybuild.rs b/actix-macros/tests/trybuild.rs
index a726d3bd..db1ccda2 100644
--- a/actix-macros/tests/trybuild.rs
+++ b/actix-macros/tests/trybuild.rs
@@ -4,6 +4,7 @@ fn compile_macros() {
     t.pass("tests/trybuild/main-01-basic.rs");
     t.compile_fail("tests/trybuild/main-02-only-async.rs");
     t.pass("tests/trybuild/main-03-fn-params.rs");
+    t.pass("tests/trybuild/main-04-system-path.rs");
 
     t.pass("tests/trybuild/test-01-basic.rs");
     t.pass("tests/trybuild/test-02-keep-attrs.rs");
diff --git a/actix-macros/tests/trybuild/main-04-system-path.rs b/actix-macros/tests/trybuild/main-04-system-path.rs
new file mode 100644
index 00000000..549dc5c7
--- /dev/null
+++ b/actix-macros/tests/trybuild/main-04-system-path.rs
@@ -0,0 +1,8 @@
+mod system {
+    pub use actix_rt::System as MySystem;
+}
+
+#[actix_rt::main(system = "system::MySystem")]
+async fn main() {
+    futures_util::future::ready(()).await
+}