config: introduce declare_event! macro

There should be no functional change here; this commit refactors
how we pre-declare CallbackSignatures with lua. The rationale is:

* It is important to correct declare single vs. multiple implementation
  event handler types before any lua code is run; we've had a couple
  of issues in the past where some part of this pre-registration
  was messed up.
* Adding a new signature requires declaring it in a global, and
  remembering to add its registration to the right place
* Declaring signatures is a bit boilerplatey and makes it hard
  grok the purpose of the event handler arguments at a glance

The declare_event! macro defined in this commit makes it a bit
more readable to declare these event types and automatically
wires up the registration to the correct spot, improving the
ergonomics significantly.

A downside of this additional layer of macro stuff that it
requires increasing the rustc recursion limit.

Cargo.lock

@@ -1112,9 +1112,11 @@ dependencies = [
  "anyhow",
  "data-encoding",
  "filenamegen",
+ "linkme",
  "metrics",
  "mlua",
  "parking_lot",
+ "paste",
  "prometheus",
  "rand 0.8.5",
  "serde",

crates/config/Cargo.toml

@@ -7,9 +7,11 @@ edition = "2021"
 anyhow = {workspace=true}
 data-encoding = {workspace=true}
 filenamegen = {workspace=true}
+linkme.workspace = true
 metrics = {workspace=true}
 mlua = {workspace=true, features=["vendored", "lua54", "async", "send", "serialize"]}
 parking_lot = {workspace=true}
+paste.workspace = true
 prometheus = {workspace=true}
 rand = {workspace=true}
 serde = {workspace=true}

crates/config/src/lib.rs

@@ -7,13 +7,14 @@ use mlua::{
     UserData, UserDataMethods, Value,
 };
 use parking_lot::FairMutex as Mutex;
+pub use paste;
 use prometheus::{CounterVec, HistogramTimer, HistogramVec};
 use serde::Serialize;
 use std::borrow::Cow;
 use std::collections::HashSet;
 use std::path::PathBuf;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::LazyLock;
+use std::sync::{LazyLock, Once};
 use std::time::Instant;
 
 pub mod epoch;
@@ -148,6 +149,8 @@ pub async fn load_config() -> anyhow::Result<LuaConfig> {
         package.set("path", path_array.join(";"))?;
     }
 
+    register_declared_events();
+
     for func in get_funcs() {
         (func)(&lua)?;
     }
@@ -642,6 +645,79 @@ where
     name: Cow<'static, str>,
 }
 
+#[linkme::distributed_slice]
+pub static CALLBACK_SIGNATURES: [fn()];
+
+/// Helper for declaring a named event handler callback signature.
+///
+/// Usage looks like:
+///
+/// ```rust,ignore
+/// declare_event! {
+/// pub static GET_Q_CONFIG_SIG: Multiple(
+///         "get_queue_config",
+///         domain: &'static str,
+///         tenant: Option<&'static str>,
+///         campaign: Option<&'static str>,
+///         routing_domain: Option<&'static str>,
+///     ) -> QueueConfig;
+/// }
+/// ```
+///
+/// A handler can be either `Single` or `Multiple`, indicating whether
+/// only a single registration or multiple registrations are permitted.
+/// The string literal is the name of the event, followed by a fn-style
+/// parameter list which names each parameter in sequence, followed by
+/// the return value.  The names are not currently used in any way,
+/// but enhance the readability of the code.
+///
+/// In addition to declaring the signature in a global, some glue
+/// is generated that will register the signature appropriately
+/// so that lua knows whether it is single or multiple and can
+/// act appropriately when `kumo.on` is called.
+#[macro_export]
+macro_rules! declare_event {
+    ($vis:vis static $sym:ident: Multiple($name:literal $(,)? $($param_name:ident: $args:ty),* $(,)? ) -> $ret:ty;) => {
+        $vis static $sym: ::std::sync::LazyLock<
+            $crate::CallbackSignature<($($args),*), $ret>> =
+                ::std::sync::LazyLock::new(|| $crate::CallbackSignature::new_with_multiple($name));
+
+        $crate::paste::paste! {
+            #[linkme::distributed_slice($crate::CALLBACK_SIGNATURES)]
+            static [<CALLBACK_SIG_REGISTER_ $sym>]: fn() = || {
+                $sym.register();
+            };
+        }
+    };
+    ($vis:vis static $sym:ident: Single($name:literal $(,)? $($param_name:ident: $args:ty),* $(,)? ) -> $ret:ty;) => {
+        $vis static $sym: ::std::sync::LazyLock<
+            $crate::CallbackSignature<($($args),*), $ret>> =
+                ::std::sync::LazyLock::new(|| $crate::CallbackSignature::new($name));
+
+        $crate::paste::paste! {
+            #[linkme::distributed_slice($crate::CALLBACK_SIGNATURES)]
+            static [<CALLBACK_SIG_REGISTER_ $sym>]: fn() = || {
+                $sym.register();
+            };
+        }
+    };
+}
+
+/// For each event handler CallbackSignature that was declared via
+/// `declare_event!`, call its `.register()` method to register
+/// it so that `kumo.on` can give appropriate messaging if misused,
+/// and so that runtime dispatch will work correctly.
+///
+/// This should be called once, prior to running any lua code.
+fn register_declared_events() {
+    static ONCE: Once = Once::new();
+    ONCE.call_once(|| {
+        for reg_func in CALLBACK_SIGNATURES {
+            reg_func();
+        }
+    });
+}
+
 impl<A, R> CallbackSignature<A, R>
 where
     A: IntoLuaMulti,

crates/kumod/src/logging/hooks.rs

@@ -2,7 +2,7 @@ use crate::logging::files::LogFileParams;
 use crate::logging::{LogCommand, LogRecordParams, LOGGING_RUNTIME};
 use crate::queue::{InsertReason, QueueManager};
 use anyhow::Context;
-use config::{load_config, CallbackSignature};
+use config::{declare_event, load_config};
 use flume::Receiver;
 pub use kumo_log_types::*;
 use kumo_template::{Template, TemplateEngine};
@@ -14,8 +14,13 @@ use std::collections::HashMap;
 use std::sync::{Arc, LazyLock};
 use tokio::sync::{Semaphore, TryAcquireError};
 
-pub static SHOULD_ENQ_LOG_RECORD_SIG: LazyLock<CallbackSignature<(Message, String), bool>> =
-    LazyLock::new(|| CallbackSignature::new_with_multiple("should_enqueue_log_record"));
+declare_event! {
+pub static SHOULD_ENQ_LOG_RECORD_SIG: Multiple(
+    "should_enqueue_log_record",
+    message: Message,
+    hook_name: String,
+) -> bool;
+}
 
 static HOOK_BACKLOG_COUNT: LazyLock<CounterVec> = LazyLock::new(|| {
     prometheus::register_counter_vec!(

crates/kumod/src/main.rs

@@ -1,7 +1,9 @@
+#![recursion_limit = "256"]
+
 use anyhow::Context;
 use chrono::Utc;
 use clap::Parser;
-use config::CallbackSignature;
+use config::{declare_event, CallbackSignature};
 use kumo_server_common::diagnostic_logging::{DiagnosticFormat, LoggingConfig};
 use kumo_server_common::start::StartConfig;
 use kumo_server_lifecycle::LifeCycle;
@@ -9,12 +11,14 @@ use nix::sys::resource::{getrlimit, setrlimit, Resource};
 use nix::unistd::{Uid, User};
 use std::path::PathBuf;
 use std::sync::atomic::{AtomicUsize, Ordering};
-use std::sync::{Arc, LazyLock};
+use std::sync::Arc;
 
-pub static PRE_INIT_SIG: LazyLock<CallbackSignature<(), ()>> =
-    LazyLock::new(|| CallbackSignature::new_with_multiple("pre_init"));
-pub static VALIDATE_SIG: LazyLock<CallbackSignature<(), ()>> =
-    LazyLock::new(|| CallbackSignature::new_with_multiple("validate_config"));
+declare_event! {
+pub static PRE_INIT_SIG: Multiple("pre_init") -> ();
+}
+declare_event! {
+pub static VALIDATE_SIG: Multiple("validate_config") -> ();
+}
 
 mod accounting;
 mod delivery_metrics;

crates/kumod/src/mod_kumo.rs

@@ -16,12 +16,6 @@ use throttle::ThrottleSpec;
 pub fn register(lua: &Lua) -> anyhow::Result<()> {
     let kumo_mod = get_or_create_module(lua, "kumo")?;
 
-    crate::queue::GET_Q_CONFIG_SIG.register();
-    crate::queue::THROTTLE_INSERT_READY_SIG.register();
-    crate::logging::hooks::SHOULD_ENQ_LOG_RECORD_SIG.register();
-    crate::PRE_INIT_SIG.register();
-    crate::VALIDATE_SIG.register();
-    crate::queue::REQUEUE_MESSAGE_SIG.register();
     crate::http_server::admin_suspend_ready_q_v1::register(lua)?;
     crate::http_server::admin_suspend_v1::register(lua)?;
     crate::http_server::admin_bounce_v1::register(lua)?;

crates/kumod/src/queue.rs

@@ -17,7 +17,7 @@ use anyhow::Context;
 use arc_swap::ArcSwap;
 use chrono::{DateTime, Utc};
 use config::epoch::{get_current_epoch, ConfigEpoch};
-use config::{load_config, CallbackSignature, LuaConfig};
+use config::{declare_event, load_config, LuaConfig};
 use crossbeam_skiplist::SkipSet;
 use dashmap::DashMap;
 use humantime::format_duration;
@@ -73,23 +73,36 @@ static TOTAL_QMAINT_RUNS: LazyLock<IntCounter> = LazyLock::new(|| {
 
 pub static QMAINT_RUNTIME: LazyLock<Runtime> =
     LazyLock::new(|| Runtime::new("qmaint", |cpus| cpus / 4, &QMAINT_THREADS).unwrap());
-pub static GET_Q_CONFIG_SIG: LazyLock<
-    CallbackSignature<
-        (
-            &'static str,
-            Option<&'static str>,
-            Option<&'static str>,
-            Option<&'static str>,
-        ),
-        QueueConfig,
-    >,
-> = LazyLock::new(|| CallbackSignature::new_with_multiple("get_queue_config"));
-pub static THROTTLE_INSERT_READY_SIG: LazyLock<CallbackSignature<Message, ()>> =
-    LazyLock::new(|| CallbackSignature::new_with_multiple("throttle_insert_ready_queue"));
-static REBIND_MESSAGE_SIG: LazyLock<CallbackSignature<(Message, HashMap<String, String>), ()>> =
-    LazyLock::new(|| CallbackSignature::new("rebind_message"));
-pub static REQUEUE_MESSAGE_SIG: LazyLock<CallbackSignature<(Message, String), ()>> =
-    LazyLock::new(|| CallbackSignature::new_with_multiple("requeue_message"));
+
+declare_event! {
+pub static GET_Q_CONFIG_SIG: Multiple(
+        "get_queue_config",
+        domain: &'static str,
+        tenant: Option<&'static str>,
+        campaign: Option<&'static str>,
+        routing_domain: Option<&'static str>,
+    ) -> QueueConfig;
+}
+declare_event! {
+pub static THROTTLE_INSERT_READY_SIG: Multiple(
+    "throttle_insert_ready_queue",
+    message: Message,
+) -> ();
+}
+declare_event! {
+static REBIND_MESSAGE_SIG: Single(
+    "rebind_message",
+    message: Message,
+    rebind_request_data: HashMap<String, String>,
+) -> ();
+}
+declare_event! {
+pub static REQUEUE_MESSAGE_SIG: Multiple(
+    "requeue_message",
+    message: Message,
+    response: String
+) -> ();
+}
 
 pub static SINGLETON_WHEEL: LazyLock<Arc<FairMutex<TimeQ<WeakMessage>>>> =
     LazyLock::new(|| Arc::new(FairMutex::new(TimeQ::new())));

crates/kumod/src/ready_queue.rs

@@ -19,7 +19,7 @@ use crate::spool::SpoolManager;
 use anyhow::Context;
 use async_trait::async_trait;
 use config::epoch::ConfigEpoch;
-use config::{load_config, CallbackSignature};
+use config::{declare_event, load_config};
 use dashmap::DashMap;
 use dns_resolver::MailExchanger;
 use kumo_api_types::egress_path::{ConfigRefreshStrategy, EgressPathConfig, MemoryReductionPolicy};
@@ -50,9 +50,15 @@ use uuid::Uuid;
 static MANAGER: LazyLock<ReadyQueueManager> = LazyLock::new(|| ReadyQueueManager::new());
 static READYQ_RUNTIME: LazyLock<Runtime> =
     LazyLock::new(|| Runtime::new("readyq", |cpus| cpus / 2, &READYQ_THREADS).unwrap());
-pub static GET_EGRESS_PATH_CONFIG_SIG: LazyLock<
-    CallbackSignature<(String, String, String), EgressPathConfig>,
-> = LazyLock::new(|| CallbackSignature::new("get_egress_path_config"));
+
+declare_event! {
+pub static GET_EGRESS_PATH_CONFIG_SIG: Single(
+    "get_egress_path_config",
+    routing_domain: String,
+    egress_source: String,
+    site_name: String
+) -> EgressPathConfig;
+}
 
 static INSERT_LATENCY: LazyLock<Histogram> = LazyLock::new(|| {
     prometheus::register_histogram!(

crates/kumod/src/smtp_server.rs

@@ -11,7 +11,7 @@ use anyhow::{anyhow, Context};
 use async_trait::async_trait;
 use chrono::Utc;
 use cidr_map::CidrSet;
-use config::{any_err, load_config, serialize_options, CallbackSignature};
+use config::{any_err, declare_event, load_config, serialize_options, CallbackSignature};
 use data_encoding::BASE64;
 use data_loader::KeySource;
 use kumo_log_types::ResolvedAddress;
@@ -47,12 +47,21 @@ use uuid::Uuid;
 
 pub const DEFERRED_QUEUE_NAME: &str = "deferred_smtp_inject.kumomta.internal";
 
-static SMTP_SERVER_MSG_RX: LazyLock<CallbackSignature<(Message, ConnectionMetaData), ()>> =
-    LazyLock::new(|| CallbackSignature::new("smtp_server_message_received"));
+declare_event! {
+static SMTP_SERVER_MSG_RX: Single(
+    "smtp_server_message_received",
+    message: Message,
+    connection_metadata: ConnectionMetaData
+) -> ();
+}
 
-static DEFERRED_SMTP_SERVER_MSG_INJECT: LazyLock<
-    CallbackSignature<(Message, ConnectionMetaData), ()>,
-> = LazyLock::new(|| CallbackSignature::new("smtp_server_message_deferred_inject"));
+declare_event! {
+static DEFERRED_SMTP_SERVER_MSG_INJECT: Single(
+    "smtp_server_message_deferred_inject",
+    message: Message,
+    connection_metadata: ConnectionMetaData
+) -> ();
+}
 
 static CRLF: LazyLock<Finder> = LazyLock::new(|| Finder::new("\r\n"));
 static TXN_LATENCY: LazyLock<Histogram> = LazyLock::new(|| {