diff --git a/crates/kumod/src/logging/disposition.rs b/crates/kumod/src/logging/disposition.rs index 1b6723a1..e39fcda8 100644 --- a/crates/kumod/src/logging/disposition.rs +++ b/crates/kumod/src/logging/disposition.rs @@ -169,7 +169,7 @@ pub async fn log_disposition(args: LogDisposition<'_>) { provider_name: provider.map(|s| s.to_string()), session_id, }; - if let Err(err) = logger.log(record).await { + if let Err(err) = logger.log(record, Some(msg.clone())).await { tracing::error!("failed to log: {err:#}"); } @@ -190,6 +190,9 @@ pub async fn log_disposition(args: LogDisposition<'_>) { .get_queue_name() .unwrap_or_else(|err| format!("{err:#}")); + let reconstructed_original_msg = None; // FIXME: try to build this from the + // parsed rfc3464 report? + for recip in &report.per_recipient { if recip.action != ReportAction::Failed { continue; @@ -261,7 +264,9 @@ pub async fn log_disposition(args: LogDisposition<'_>) { session_id, }; - if let Err(err) = logger.log(record).await { + if let Err(err) = + logger.log(record, reconstructed_original_msg.clone()).await + { tracing::error!("failed to log: {err:#}"); } } diff --git a/crates/kumod/src/logging/disposition_hooks.rs b/crates/kumod/src/logging/disposition_hooks.rs new file mode 100644 index 00000000..6c8ab076 --- /dev/null +++ b/crates/kumod/src/logging/disposition_hooks.rs @@ -0,0 +1,47 @@ +use crate::logging::LogRecordParams; +use config::{load_config, CallbackSignature}; +use kumo_log_types::{JsonLogRecord, RecordType}; +use message::Message; +use mlua::{IntoLua, LuaSerdeExt}; +use serde::Deserialize; +use std::collections::HashMap; + +#[derive(Clone)] +pub struct RecordWrapper(JsonLogRecord); + +impl IntoLua for RecordWrapper { + fn into_lua(self, lua: &mlua::Lua) -> mlua::Result { + lua.to_value(&self.0) + } +} + +#[derive(Deserialize, Clone, Debug)] +#[serde(deny_unknown_fields)] +pub struct DispHookParams { + /// The unique name to identify this instance of the log hook + pub name: String, + + #[serde(default)] + pub per_record: HashMap, +} + +impl DispHookParams { + pub async fn do_record( + sig: &CallbackSignature<(Message, RecordWrapper), ()>, + msg: Message, + record: JsonLogRecord, + ) -> anyhow::Result<()> { + tracing::trace!("do_record {record:?}"); + + if record.reception_protocol.as_deref() == Some("LogRecord") { + return Ok(()); + } + + let mut lua_config = load_config().await?; + lua_config + .async_call_callback(&sig, (msg.clone(), RecordWrapper(record))) + .await?; + lua_config.put(); + anyhow::Result::<()>::Ok(()) + } +} diff --git a/crates/kumod/src/logging/files.rs b/crates/kumod/src/logging/files.rs index cf45a019..3756ff7c 100644 --- a/crates/kumod/src/logging/files.rs +++ b/crates/kumod/src/logging/files.rs @@ -204,7 +204,7 @@ impl LogThreadState { tracing::debug!("LogCommand::Terminate received. Stopping writing logs"); break; } - LogCommand::Record(record) => { + LogCommand::Record(record, _msg) => { if let Err(err) = self.do_record(record) { tracing::error!("failed to log: {err:#}"); }; diff --git a/crates/kumod/src/logging/hooks.rs b/crates/kumod/src/logging/hooks.rs index ee9cbd1d..f080e0f1 100644 --- a/crates/kumod/src/logging/hooks.rs +++ b/crates/kumod/src/logging/hooks.rs @@ -97,7 +97,7 @@ impl LogHookState { tracing::debug!("LogCommand::Terminate received. Stopping writing logs"); break; } - LogCommand::Record(record) => { + LogCommand::Record(record, _msg) => { if let Err(err) = self.do_record(record).await { tracing::error!("failed to log: {err:#}"); }; diff --git a/crates/kumod/src/logging/mod.rs b/crates/kumod/src/logging/mod.rs index a6d94d3a..4af817d0 100644 --- a/crates/kumod/src/logging/mod.rs +++ b/crates/kumod/src/logging/mod.rs @@ -1,8 +1,9 @@ use crate::logging::classify::{apply_classification, ClassifierParams}; +use crate::logging::disposition_hooks::{DispHookParams, RecordWrapper}; use crate::logging::files::{LogFileParams, LogThreadState}; use crate::logging::hooks::{LogHookParams, LogHookState}; use anyhow::Context; -use config::{any_err, from_lua_value, get_or_create_module}; +use config::{any_err, from_lua_value, get_or_create_module, CallbackSignature}; use flume::{bounded, Sender, TrySendError}; pub use kumo_log_types::*; use kumo_server_common::disk_space::MonitoredPath; @@ -23,6 +24,7 @@ use tokio::task::JoinHandle; pub(crate) mod classify; pub(crate) mod disposition; +pub(crate) mod disposition_hooks; pub(crate) mod files; pub(crate) mod hooks; pub(crate) mod rejection; @@ -82,19 +84,32 @@ fn default_true() -> bool { #[derive(Debug)] pub(crate) enum LogCommand { - Record(JsonLogRecord), + Record(JsonLogRecord, Option), Terminate, } +enum LoggerImpl { + /// Queued to another thread to dispatch + Queue { + sender: Sender, + thread: TokioMutex>>, + }, + /// Processed immediate in the context of log_disposition. + /// Necessary for things that might operate on (eg: read) the + /// originating message prior to it being removed from the spool. + /// We can't queue those operations because the caller might remove + /// the message from the spool before the log event would be + /// dispatched. + Immediate(Arc>), +} + pub struct Logger { - sender: Sender, - thread: TokioMutex>>, + implementation: LoggerImpl, meta: Vec, headers: Vec, enabled: HashMap, filter_event: Option, hook_name: Option, - #[allow(unused)] name: String, submit_latency: Histogram, } @@ -104,6 +119,46 @@ impl Logger { LOGGER.lock().iter().map(Arc::clone).collect() } + pub async fn init_disp_hook(params: DispHookParams) -> anyhow::Result<()> { + let mut loggers = LOGGER.lock(); + + if loggers + .iter() + .any(|existing| existing.hook_name.as_deref() == Some(params.name.as_str())) + { + anyhow::bail!( + "A logging hook with name `{}` has already been registered", + params.name + ); + } + + let mut enabled = HashMap::new(); + for (kind, cfg) in ¶ms.per_record { + enabled.insert(*kind, cfg.enable); + } + + let hook_name = params.name.to_string(); + let name = format!("hook-{hook_name}"); + + let sig = CallbackSignature::new(format!("log_disposition_{hook_name}")); + + let submit_latency = SUBMIT_LATENCY.get_metric_with_label_values(&[&name])?; + + let logger = Self { + implementation: LoggerImpl::Immediate(Arc::new(sig)), + meta: Default::default(), + headers: Default::default(), + enabled, + filter_event: None, + hook_name: Some(hook_name), + name, + submit_latency, + }; + + loggers.push(Arc::new(logger)); + Ok(()) + } + pub async fn init_hook(params: LogHookParams) -> anyhow::Result<()> { let mut loggers = LOGGER.lock(); @@ -150,10 +205,13 @@ impl Logger { })?; let submit_latency = SUBMIT_LATENCY.get_metric_with_label_values(&[&name])?; - - let logger = Self { + let implementation = LoggerImpl::Queue { sender, thread: TokioMutex::new(Some(thread)), + }; + + let logger = Self { + implementation, meta, headers, enabled, @@ -217,9 +275,13 @@ impl Logger { let submit_latency = SUBMIT_LATENCY.get_metric_with_label_values(&[&name])?; - let logger = Self { + let implementation = LoggerImpl::Queue { sender, thread: TokioMutex::new(Some(thread)), + }; + + let logger = Self { + implementation, meta, headers, enabled, @@ -243,34 +305,54 @@ impl Logger { true } - pub async fn log(&self, mut record: JsonLogRecord) -> anyhow::Result<()> { + pub async fn log(&self, mut record: JsonLogRecord, msg: Option) -> anyhow::Result<()> { let _timer = self.submit_latency.start_timer(); apply_classification(&mut record).await; - match self.sender.try_send(LogCommand::Record(record)) { - Ok(()) => Ok(()), - Err(TrySendError::Full(record)) => { - SUBMIT_FULL - .get_metric_with_label_values(&[&self.name]) - .expect("get counter") - .inc(); - self.sender.send_async(record).await?; - Ok(()) + match &self.implementation { + LoggerImpl::Immediate(sig) => { + match msg { + Some(msg) => { + // Need to put this future on the heap, otherwise we can consume + // too much stack + let future = Box::pin(DispHookParams::do_record(sig, msg, record)); + future.await + } + None => Ok(()), + } + } + LoggerImpl::Queue { sender, .. } => { + match sender.try_send(LogCommand::Record(record, msg)) { + Ok(()) => Ok(()), + Err(TrySendError::Full(record)) => { + SUBMIT_FULL + .get_metric_with_label_values(&[&self.name]) + .expect("get counter") + .inc(); + sender.send_async(record).await?; + Ok(()) + } + Err(TrySendError::Disconnected(_)) => anyhow::bail!("log channel was closed"), + } } - Err(TrySendError::Disconnected(_)) => anyhow::bail!("log channel was closed"), } } pub async fn signal_shutdown() { let loggers = Self::get_loggers(); for logger in loggers.iter() { - tracing::debug!("Terminating a logger"); - logger.sender.send_async(LogCommand::Terminate).await.ok(); - tracing::debug!("Joining that logger"); - let res = match logger.thread.lock().await.take() { - Some(task) => Some(task.await), - None => None, - }; - tracing::debug!("Joined -> {res:?}"); + match &logger.implementation { + LoggerImpl::Immediate(_) => {} + LoggerImpl::Queue { sender, thread } => { + tracing::debug!("Terminating a logger"); + sender.send_async(LogCommand::Terminate).await.ok(); + tracing::debug!("Joining that logger"); + let res = match thread.lock().await.take() { + Some(task) => Some(task.await), + None => None, + }; + tracing::debug!("Joined -> {res:?}"); + } + } } } @@ -377,5 +459,13 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { })?, )?; + kumo_mod.set( + "configure_log_disposition_hook", + lua.create_async_function(|lua, params: LuaValue| async move { + let params: DispHookParams = from_lua_value(&lua, params)?; + Logger::init_disp_hook(params).await.map_err(any_err) + })?, + )?; + Ok(()) } diff --git a/crates/kumod/src/logging/rejection.rs b/crates/kumod/src/logging/rejection.rs index 463cc8ef..5a79f5ee 100644 --- a/crates/kumod/src/logging/rejection.rs +++ b/crates/kumod/src/logging/rejection.rs @@ -61,7 +61,7 @@ pub async fn log_rejection(args: LogRejection) { provider_name: None, session_id: args.session_id, }; - if let Err(err) = logger.log(record).await { + if let Err(err) = logger.log(record, None).await { tracing::error!("failed to log: {err:#}"); } }