diff --git a/core/codegen/src/attribute/entry/launch.rs b/core/codegen/src/attribute/entry/launch.rs index a97c5eaa5d..48dcda3d14 100644 --- a/core/codegen/src/attribute/entry/launch.rs +++ b/core/codegen/src/attribute/entry/launch.rs @@ -113,7 +113,9 @@ impl EntryAttr for Launch { #[allow(dead_code)] #f #vis #sig { - #_error::Error::report(::rocket::async_main(#launch)) + ::rocket::async_main(async move { + #_error::Error::report(#launch.await) + }) } )) } diff --git a/core/lib/tests/drop-in-async-context.rs b/core/lib/tests/drop-in-async-context.rs new file mode 100644 index 0000000000..d774f5e308 --- /dev/null +++ b/core/lib/tests/drop-in-async-context.rs @@ -0,0 +1,67 @@ +#[macro_use] +extern crate rocket; + +use rocket::{Build, Config, Rocket}; +use rocket::fairing::AdHoc; +use rocket::figment::Figment; + +struct AsyncDropInAsync; + +impl Drop for AsyncDropInAsync { + fn drop(&mut self) { + // Ensure that managed state is dropped inside of an async context by + // ensuring that we do not panic when fetching the current runtime. + // + // Crates like rocket_sync_db_pools spawn tasks to asynchronously + // complete pool shutdown which must be done in an async context or else + // the spawn will panic. We want to ensure that does not happen. + let _ = rocket::tokio::runtime::Handle::current(); + } +} + +fn rocket() -> Rocket { + let figment = Figment::from(Config::debug_default()) + .merge(("address", "tcp:127.0.0.1:0")); + + rocket::custom(figment) + .manage(AsyncDropInAsync) + .attach(AdHoc::on_liftoff("Shutdown", |rocket| Box::pin(async { + rocket.shutdown().notify(); + }))) +} + +mod launch { + #[launch] + fn launch() -> _ { + super::rocket() + } + + #[test] + fn test_launch() { + main(); + } +} + +mod main { + #[rocket::main] + async fn main() { + super::rocket().launch().await.unwrap(); + } + + #[test] + fn test_main() { + main(); + } + + #[test] + fn test_execute() { + rocket::execute(async { + super::rocket().launch().await.unwrap(); + }); + } + + #[test] + fn test_execute_directly() { + rocket::execute(super::rocket().launch()).unwrap(); + } +}