diff --git a/Cargo.lock b/Cargo.lock index 0d3a85ce..f407066a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1112,9 +1112,11 @@ dependencies = [ "anyhow", "data-encoding", "filenamegen", + "linkme", "metrics", "mlua", "parking_lot", + "paste", "prometheus", "rand 0.8.5", "serde", diff --git a/crates/config/Cargo.toml b/crates/config/Cargo.toml index 64d15c59..22b108bc 100644 --- a/crates/config/Cargo.toml +++ b/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} diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 1acec53d..4d6eed07 100644 --- a/crates/config/src/lib.rs +++ b/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 { 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 []: 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 []: 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 CallbackSignature where A: IntoLuaMulti, diff --git a/crates/kumod/src/logging/hooks.rs b/crates/kumod/src/logging/hooks.rs index 87ee6e67..3511fbae 100644 --- a/crates/kumod/src/logging/hooks.rs +++ b/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> = - 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 = LazyLock::new(|| { prometheus::register_counter_vec!( diff --git a/crates/kumod/src/main.rs b/crates/kumod/src/main.rs index 808269ab..5ac39af7 100644 --- a/crates/kumod/src/main.rs +++ b/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> = - LazyLock::new(|| CallbackSignature::new_with_multiple("pre_init")); -pub static VALIDATE_SIG: LazyLock> = - 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; diff --git a/crates/kumod/src/mod_kumo.rs b/crates/kumod/src/mod_kumo.rs index c24b060e..f137d5e1 100644 --- a/crates/kumod/src/mod_kumo.rs +++ b/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)?; diff --git a/crates/kumod/src/queue.rs b/crates/kumod/src/queue.rs index b7c386f3..345e1eb9 100644 --- a/crates/kumod/src/queue.rs +++ b/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 = LazyLock::new(|| { pub static QMAINT_RUNTIME: LazyLock = 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> = - LazyLock::new(|| CallbackSignature::new_with_multiple("throttle_insert_ready_queue")); -static REBIND_MESSAGE_SIG: LazyLock), ()>> = - LazyLock::new(|| CallbackSignature::new("rebind_message")); -pub static REQUEUE_MESSAGE_SIG: LazyLock> = - 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, +) -> (); +} +declare_event! { +pub static REQUEUE_MESSAGE_SIG: Multiple( + "requeue_message", + message: Message, + response: String +) -> (); +} pub static SINGLETON_WHEEL: LazyLock>>> = LazyLock::new(|| Arc::new(FairMutex::new(TimeQ::new()))); diff --git a/crates/kumod/src/ready_queue.rs b/crates/kumod/src/ready_queue.rs index 088ddc7d..08924502 100644 --- a/crates/kumod/src/ready_queue.rs +++ b/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 = LazyLock::new(|| ReadyQueueManager::new()); static READYQ_RUNTIME: LazyLock = 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 = LazyLock::new(|| { prometheus::register_histogram!( diff --git a/crates/kumod/src/smtp_server.rs b/crates/kumod/src/smtp_server.rs index e6b277d4..db9a604a 100644 --- a/crates/kumod/src/smtp_server.rs +++ b/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> = - 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 = LazyLock::new(|| Finder::new("\r\n")); static TXN_LATENCY: LazyLock = LazyLock::new(|| {