From 2a478bcadd783cfbdf284d356f86f526ad94eda5 Mon Sep 17 00:00:00 2001
From: Steven Fackler <sfackler@palantir.com>
Date: Fri, 10 Oct 2025 14:21:18 -0400
Subject: [PATCH] Support async init functions

---
 witchcraft-server-macros/src/lib.rs | 50 ++++++++++++++++++++++-
 witchcraft-server/src/lib.rs        | 62 ++++++++++++++++++++++++++++-
 2 files changed, 109 insertions(+), 3 deletions(-)

diff --git a/witchcraft-server-macros/src/lib.rs b/witchcraft-server-macros/src/lib.rs
index 8966072b..55225593 100644
--- a/witchcraft-server-macros/src/lib.rs
+++ b/witchcraft-server-macros/src/lib.rs
@@ -65,6 +65,49 @@ use syn::{Error, ItemFn};
 ///     witchcraft_server::init(inner_main)
 /// }
 /// ```
+///
+/// Async functions are also supported:
+///
+/// ```ignore
+/// use conjure_error::Error;
+/// use witchcraft_server::config::install::InstallConfig;
+/// use witchcraft_server::config::runtime::RuntimeConfig;
+/// use witchcraft_server::Witchcraft;
+/// use refreshable::Refreshable;
+///
+/// #[witchcraft_server::main]
+/// async fn main(
+///     install: InstallConfig,
+///     runtime: Refreshable<RuntimeConfig, Error>,
+///     wc: &mut Witchcraft,
+/// ) -> Result<(), Error> {
+///     // initialization code...
+///     Ok(())
+/// }
+/// ```
+///
+/// Expands to:
+///
+/// ```ignore
+/// use conjure_error::Error;
+/// use witchcraft_server::config::install::InstallConfig;
+/// use witchcraft_server::config::runtime::RuntimeConfig;
+/// use witchcraft_server::Witchcraft;
+/// use refreshable::Refreshable;
+///
+/// fn main() {
+///     async fn inner_main(
+///         install: InstallConfig,
+///         runtime: Refreshable<RuntimeConfig, Error>,
+///         wc: &mut Witchcraft,
+///     ) -> Result<(), Error> {
+///         // initialization code...
+///         Ok(())
+///     }
+///
+///     witchcraft_server::init_async(inner_main)
+/// }
+/// ```
 #[proc_macro_attribute]
 pub fn main(args: TokenStream, input: TokenStream) -> TokenStream {
     if !args.is_empty() {
@@ -84,11 +127,16 @@ pub fn main(args: TokenStream, input: TokenStream) -> TokenStream {
     let vis = &function.vis;
     let name = &function.sig.ident;
 
+    let init = match function.sig.asyncness {
+        Some(_) => quote!(init_async),
+        None => quote!(init),
+    };
+
     quote! {
         #vis fn #name() {
             #function
 
-            witchcraft_server::init(#name)
+            witchcraft_server::#init(#name)
         }
     }
     .into()
diff --git a/witchcraft-server/src/lib.rs b/witchcraft-server/src/lib.rs
index f1dc60ab..2e80ce7f 100644
--- a/witchcraft-server/src/lib.rs
+++ b/witchcraft-server/src/lib.rs
@@ -17,18 +17,20 @@
 //!
 //! The entrypoint of a Witchcraft server is an initialization function annotated with `#[witchcraft_server::main]`:
 //!
-//! ```ignore
+//! ```no_run
 //! use conjure_error::Error;
 //! use refreshable::Refreshable;
+//! use witchcraft_server::Witchcraft;
 //! use witchcraft_server::config::install::InstallConfig;
 //! use witchcraft_server::config::runtime::RuntimeConfig;
 //!
 //! #[witchcraft_server::main]
 //! fn main(
 //!     install: InstallConfig,
-//!     runtime: Refreshable<RuntimeConfig>,
+//!     runtime: Refreshable<RuntimeConfig, Error>,
 //!     wc: &mut Witchcraft,
 //! ) -> Result<(), Error> {
+//!     # #[cfg(any())]
 //!     wc.api(CustomApiEndpoints::new(CustomApiResource));
 //!
 //!     Ok(())
@@ -38,6 +40,28 @@
 //! The function is provided with the server's install and runtime configuration, as well as the [`Witchcraft`] object
 //! which can be used to configure the server. Once the initialization function returns, the server will start.
 //!
+//! Async init functions are also supported:
+//!
+//! ```no_run
+//! use conjure_error::Error;
+//! use refreshable::Refreshable;
+//! use witchcraft_server::Witchcraft;
+//! use witchcraft_server::config::install::InstallConfig;
+//! use witchcraft_server::config::runtime::RuntimeConfig;
+//!
+//! #[witchcraft_server::main]
+//! async fn main(
+//!     install: InstallConfig,
+//!     runtime: Refreshable<RuntimeConfig, Error>,
+//!     wc: &mut Witchcraft,
+//! ) -> Result<(), Error> {
+//!     # #[cfg(any())]
+//!     wc.api(CustomApiEndpoints::new(CustomApiResource));
+//!
+//!     Ok(())
+//! }
+//! ```
+//!
 //! ## Note
 //!
 //! The initialization function is expected to return quickly - any long-running work required should happen in the
@@ -374,6 +398,18 @@ where
     init_with_configs(init, configs::load_install::<I>, configs::load_runtime::<R>)
 }
 
+/// Initializes a Witchcraft server.
+///
+/// This is equivalent to [`init`], except that it takes an asynchronous init function.
+pub fn init_async<I, R, F>(init_async: F)
+where
+    I: AsRef<InstallConfig> + DeserializeOwned,
+    R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
+    F: AsyncFnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error>,
+{
+    init(|install, runtime, wc| wc.handle.clone().block_on(init_async(install, runtime, wc)))
+}
+
 /// Initializes a Witchcraft server with custom config loaders.
 ///
 /// `init` is invoked with the install and runtime configs from the provided loaders as well as the [`Witchcraft`]
@@ -401,6 +437,28 @@ where
     process::exit(ret);
 }
 
+/// Initializes a Witchcraft server with custom config loaders.
+///
+/// This is equivalent to [`init_with_configs`] except that it takes an asynchrnous init function.
+pub fn init_with_configs_async<I, R, F, LI, LR>(init_async: F, load_install: LI, load_runtime: LR)
+where
+    I: AsRef<InstallConfig> + DeserializeOwned,
+    R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
+    F: AsyncFnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error>,
+    LI: FnOnce() -> Result<I, Error>,
+    LR: FnOnce(&Handle, &Arc<AtomicBool>) -> Result<Refreshable<R, Error>, Error>,
+{
+    init_with_configs(
+        |install, runtime, wc| {
+            wc.handle()
+                .clone()
+                .block_on(init_async(install, runtime, wc))
+        },
+        load_install,
+        load_runtime,
+    );
+}
+
 fn init_inner<I, R, F, LI, LR>(
     init: F,
     load_install: LI,