Add a witchcraft builder and make runtime configurable

witchcraft-server-macros/Cargo.toml

@@ -13,4 +13,5 @@ proc-macro = true
 [dependencies]
 proc-macro2 = "1"
 quote = "1"
+structmeta = "0.3.0"
 syn = { version = "2", features = ["full"] }

witchcraft-server-macros/src/lib.rs

@@ -13,16 +13,22 @@
 // limitations under the License.
 #![allow(clippy::needless_doctest_main)]
 use proc_macro::TokenStream;
-use proc_macro2::Span;
 use quote::quote;
-use syn::{Error, ItemFn};
+use structmeta::StructMeta;
+use syn::{Error, Expr, ItemFn};
 
 /// Marks the entrypoint of a Witchcraft server.
 ///
 /// This macro should be applied to a function taking 3 arguments: the server's install config, a `Refreshable` of the
 /// server's runtime config, and a mutable reference to the `Witchcraft` context object. It is a simple convenience
 /// function that wraps the annotated function in one that passes it to the `witchcraft_server::init` function.
 ///
+/// # Parameters
+///
+/// The attribute allows configuration of the Witchcraft builder via optional parameters:
+///
+/// * `conjure_runtime` - An expression which creates the `Arc<ConjureRuntime>` used for application endpoints.
+///
 /// # Examples
 ///
 /// ```ignore
@@ -62,20 +68,16 @@ use syn::{Error, ItemFn};
 ///         Ok(())
 ///     }
 ///
-///     witchcraft_server::init(inner_main)
+///     witchcraft_server::Builder::new()
+///         .init(inner_main)
 /// }
 /// ```
 #[proc_macro_attribute]
 pub fn main(args: TokenStream, input: TokenStream) -> TokenStream {
-    if !args.is_empty() {
-        return with_error(
-            input,
-            Error::new(
-                Span::call_site(),
-                "#[witchcraft_server::main] does not take arguments",
-            ),
-        );
-    }
+    let args = match syn::parse::<MainArgs>(args) {
+        Ok(args) => args,
+        Err(e) => return with_error(input, e),
+    };
 
     let function = match syn::parse::<ItemFn>(input.clone()) {
         Ok(function) => function,
@@ -84,16 +86,25 @@ pub fn main(args: TokenStream, input: TokenStream) -> TokenStream {
     let vis = &function.vis;
     let name = &function.sig.ident;
 
+    let conjure_runtime_call = args.conjure_runtime.map(|e| quote!(.conjure_runtime(#e)));
+
     quote! {
         #vis fn #name() {
             #function
 
-            witchcraft_server::init(#name)
+            witchcraft_server::Builder::new()
+                #conjure_runtime_call
+                .init(#name)
         }
     }
     .into()
 }
 
+#[derive(StructMeta)]
+struct MainArgs {
+    conjure_runtime: Option<Expr>,
+}
+
 fn with_error(mut tokens: TokenStream, error: Error) -> TokenStream {
     tokens.extend(TokenStream::from(error.into_compile_error()));
     tokens

witchcraft-server/src/lib.rs

@@ -360,6 +360,75 @@ mod witchcraft;
 #[global_allocator]
 static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
+/// A builder to configure a [`Witchcraft`].
+#[derive(Default)]
+pub struct Builder {
+    conjure_runtime: Arc<ConjureRuntime>,
+}
+
+impl Builder {
+    /// Creates a new builder with default settings.
+    #[inline]
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Sets the [`ConjureRuntime`] used for application endpoints.
+    ///
+    /// This does not affect the runtime used by the server's management endpoints.
+    #[inline]
+    pub fn conjure_runtime(mut self, conjure_runtime: Arc<ConjureRuntime>) -> Self {
+        self.conjure_runtime = conjure_runtime;
+        self
+    }
+
+    /// Initializes a Witchcraft server.
+    ///
+    /// `init` is invoked with the parsed install and runtime configs as well as the [`Witchcraft`] context object. It
+    /// is expected to return quickly; any long running initialization should be spawned off into the background to run
+    /// asynchronously.
+    pub fn init<I, R, F>(self, init: F)
+    where
+        I: AsRef<InstallConfig> + DeserializeOwned,
+        R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
+        F: FnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error>,
+    {
+        self.init_with_configs(init, configs::load_install::<I>, configs::load_runtime::<R>)
+    }
+
+    /// 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`]
+    /// context object. It is expected to return quickly; any long running initialization should be spawned off into
+    /// the background to run asynchronously.
+    pub fn init_with_configs<I, R, F, LI, LR>(
+        self,
+        init: F,
+        load_install: LI,
+        load_runtime: LR,
+    ) -> !
+    where
+        I: AsRef<InstallConfig> + DeserializeOwned,
+        R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
+        F: FnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error>,
+        LI: FnOnce() -> Result<I, Error>,
+        LR: FnOnce(&Handle, &Arc<AtomicBool>) -> Result<Refreshable<R, Error>, Error>,
+    {
+        let mut runtime_guard = None;
+
+        let ret = match init_inner(self, init, load_install, load_runtime, &mut runtime_guard) {
+            Ok(()) => 0,
+            Err(e) => {
+                fatal!("error starting server", error: e);
+                1
+            }
+        };
+        drop(runtime_guard);
+
+        process::exit(ret);
+    }
+}
+
 /// Initializes a Witchcraft server.
 ///
 /// `init` is invoked with the parsed install and runtime configs as well as the [`Witchcraft`] context object. It
@@ -371,7 +440,7 @@ where
     R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
     F: FnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error>,
 {
-    init_with_configs(init, configs::load_install::<I>, configs::load_runtime::<R>)
+    Builder::new().init(init)
 }
 
 /// Initializes a Witchcraft server with custom config loaders.
@@ -387,21 +456,11 @@ where
     LI: FnOnce() -> Result<I, Error>,
     LR: FnOnce(&Handle, &Arc<AtomicBool>) -> Result<Refreshable<R, Error>, Error>,
 {
-    let mut runtime_guard = None;
-
-    let ret = match init_inner(init, load_install, load_runtime, &mut runtime_guard) {
-        Ok(()) => 0,
-        Err(e) => {
-            fatal!("error starting server", error: e);
-            1
-        }
-    };
-    drop(runtime_guard);
-
-    process::exit(ret);
+    Builder::new().init_with_configs(init, load_install, load_runtime)
 }
 
 fn init_inner<I, R, F, LI, LR>(
+    builder: Builder,
     init: F,
     load_install: LI,
     load_runtime: LR,
@@ -512,6 +571,7 @@ where
         thread_pool: None,
         endpoints: vec![],
         shutdown_hooks: ShutdownHooks::new(),
+        // intentionally using a default runtime for the management endpoints
         conjure_runtime: Arc::new(ConjureRuntime::new()),
     };
 
@@ -543,6 +603,8 @@ where
         }
     }
 
+    witchcraft.conjure_runtime = builder.conjure_runtime;
+
     init(install_config, runtime_config, &mut witchcraft)?;
 
     witchcraft

witchcraft-server/src/witchcraft.rs

@@ -11,6 +11,7 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+
 use crate::blocking::conjure::ConjureBlockingEndpoint;
 use crate::blocking::pool::ThreadPool;
 use crate::debug::DiagnosticRegistry;