From 2bd7410ca64052e8792adfc5f5b7d23dc41eed7b Mon Sep 17 00:00:00 2001 From: Wez Furlong Date: Sat, 1 Mar 2025 08:51:00 -0700 Subject: [PATCH] config: fix cancellation safety issue with async lua calls This commit addresses a tricky interaction with lua's garbage collection and cancellation safety of async lua call handling. In Rust, when a future is cancelled it is simply not called any more, and any resources it owns will be released when it is dropped. In mlua's Rust bindings, async calls are futures that are wrapped up into state whose ownership is managed in the lua context stack, which in turn has managed via its own garbage collection process. In KumoMTA, we implement a pool of Lua contexts which we keep around for reuse. Our logic will return our LuaConfig object back to this pool automatically when it is dropped. In a situation where tokio timeouts or select! are used with a future that transitively contains a lua async call and that call is cancelled, that can leave async locks/semaphores owned in the un-collected objects in the lua stack. For async locks/semaphores that were in the locked state, this can leave them locked until the owning lua context is expired from the pool, either due to age, capacity or use count being hit. That's bad because it can block certain processing flows for 5 minutes (the default `kumo.set_max_lua_context_age`). What this commit does is switch away from automatically returning the LuaConfig to the pool on Drop. In that case the lua context is destroy on Drop, immediately allowing destructors to release any/all owned resources. We do still want to benefit from the pool, so the bulk of this commit is changing the consumers of load_config to explicitly call `config.put()` to put the object back into the pool, but only in the success case. This does penalize flows that grab a config and then `kumo.reject` as their final status. At the present time, the majority of kumomta deployments are very much focused on (relatively) trusted senders and an outbound use case, so this isn't a terrible situation for the time being. --- crates/config/src/lib.rs | 34 ++++++++++++++----- .../src/http_server/auth.rs | 14 ++++---- crates/kumo-server-common/src/lib.rs | 6 +++- crates/kumod/src/http_server/inject_v1.rs | 2 ++ crates/kumod/src/logging/disposition.rs | 5 ++- crates/kumod/src/logging/hooks.rs | 1 + crates/kumod/src/main.rs | 2 ++ crates/kumod/src/mod_kumo.rs | 1 + crates/kumod/src/queue.rs | 5 +++ crates/kumod/src/ready_queue.rs | 1 + crates/kumod/src/smtp_dispatcher.rs | 3 ++ crates/kumod/src/smtp_server.rs | 17 ++++++---- crates/tsa-daemon/src/main.rs | 1 + crates/tsa-daemon/src/shaping_config.rs | 1 + 14 files changed, 70 insertions(+), 23 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index b7aa3e95..1acec53d 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -87,17 +87,10 @@ pub struct LuaConfig { inner: Option, } -impl Drop for LuaConfig { - fn drop(&mut self) { - if let Some(inner) = self.inner.take() { - pool_put(inner); - } - } -} - pub async fn set_policy_path(path: PathBuf) -> anyhow::Result<()> { POLICY_FILE.lock().replace(path); - load_config().await?; + let config = load_config().await?; + config.put(); Ok(()) } @@ -223,6 +216,29 @@ impl LuaConfig { } } + /// Explicitly put the config object back into its containing pool. + /// Ideally we'd do this automatically when the object is dropped, + /// but lua's garbage collection makes this problematic: + /// if a future whose graph contains an async lua call within + /// this config object is cancelled (eg: simply stopped without + /// calling it again), and the config object is not explicitly garbage + /// collected, any futures and data owned by any dependencies of + /// the cancelled future remain alive until the next gc run, + /// which can cause things like async locks and semaphores to + /// have a lifetime extended by the maximum age of the lua context. + /// + /// The combat this, consumers of LuaConfig should explicitly + /// call `config.put()` after successfully using the config + /// object. + /// + /// Or framing it another way: consumers must not call `config.put()` + /// if a transitive dep might have been cancelled. + pub fn put(mut self) { + if let Some(inner) = self.inner.take() { + pool_put(inner); + } + } + pub async fn async_call_callback( &mut self, sig: &CallbackSignature, diff --git a/crates/kumo-server-common/src/http_server/auth.rs b/crates/kumo-server-common/src/http_server/auth.rs index 06ae8b1d..527da708 100644 --- a/crates/kumo-server-common/src/http_server/auth.rs +++ b/crates/kumo-server-common/src/http_server/auth.rs @@ -51,22 +51,24 @@ impl AuthKind { async fn validate_impl(&self) -> anyhow::Result { let mut config = load_config().await?; - match self { - Self::TrustedIp(_) => Ok(true), + let result = match self { + Self::TrustedIp(_) => true, Self::Basic { user, password } => { let sig = CallbackSignature::<(String, Option), bool>::new( "http_server_validate_auth_basic", ); - Ok(config + config .async_call_callback(&sig, (user.to_string(), password.clone())) - .await?) + .await? } Self::Bearer { token } => { let sig = CallbackSignature::::new("http_server_validate_auth_bearer"); - Ok(config.async_call_callback(&sig, token.to_string()).await?) + config.async_call_callback(&sig, token.to_string()).await? } - } + }; + config.put(); + Ok(result) } async fn lookup_cache(&self) -> Option> { diff --git a/crates/kumo-server-common/src/lib.rs b/crates/kumo-server-common/src/lib.rs index 6a0c58bb..ec202bc4 100644 --- a/crates/kumo-server-common/src/lib.rs +++ b/crates/kumo-server-common/src/lib.rs @@ -325,7 +325,11 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { config .convert_args_and_call_callback(&sig, &self.args) - .await + .await?; + + config.put(); + + Ok(()) } } diff --git a/crates/kumod/src/http_server/inject_v1.rs b/crates/kumod/src/http_server/inject_v1.rs index 8a0fab91..95a601ab 100644 --- a/crates/kumod/src/http_server/inject_v1.rs +++ b/crates/kumod/src/http_server/inject_v1.rs @@ -695,6 +695,8 @@ async fn inject_v1_impl( } } } + config.put(); + Ok(Json(InjectV1Response { success_count, fail_count, diff --git a/crates/kumod/src/logging/disposition.rs b/crates/kumod/src/logging/disposition.rs index 77ec5ec8..08c4621e 100644 --- a/crates/kumod/src/logging/disposition.rs +++ b/crates/kumod/src/logging/disposition.rs @@ -78,7 +78,10 @@ pub async fn log_disposition(args: LogDisposition<'_>) { let enqueue: bool = match lua_config.async_call_callback(&log_sig, msg.clone()).await { - Ok(b) => b, + Ok(b) => { + lua_config.put(); + b + } Err(err) => { tracing::error!( "error while calling {name} event for log filter: {err:#}" diff --git a/crates/kumod/src/logging/hooks.rs b/crates/kumod/src/logging/hooks.rs index c79c6a43..890f4fe6 100644 --- a/crates/kumod/src/logging/hooks.rs +++ b/crates/kumod/src/logging/hooks.rs @@ -170,6 +170,7 @@ impl LogHookState { let enqueue: bool = lua_config .async_call_callback(&SHOULD_ENQ_LOG_RECORD_SIG, (msg.clone(), name)) .await?; + lua_config.put(); if enqueue { let queue_name = msg.get_queue_name()?; diff --git a/crates/kumod/src/main.rs b/crates/kumod/src/main.rs index 6ccda133..3943563d 100644 --- a/crates/kumod/src/main.rs +++ b/crates/kumod/src/main.rs @@ -293,6 +293,8 @@ async fn perform_init(opts: Opt) -> anyhow::Result<()> { .context("start_spool")?; } + config.put(); + Ok(()) } diff --git a/crates/kumod/src/mod_kumo.rs b/crates/kumod/src/mod_kumo.rs index 967c6a22..c24b060e 100644 --- a/crates/kumod/src/mod_kumo.rs +++ b/crates/kumod/src/mod_kumo.rs @@ -159,6 +159,7 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { crate::queue::Queue::call_get_queue_config(&queue_name, &mut config) .await .map_err(any_err)?; + config.put(); lua.to_value(&queue_config) })?, )?; diff --git a/crates/kumod/src/queue.rs b/crates/kumod/src/queue.rs index d93b5a67..7edf6d6b 100644 --- a/crates/kumod/src/queue.rs +++ b/crates/kumod/src/queue.rs @@ -1090,6 +1090,8 @@ impl Queue { let queue_config = Self::call_get_queue_config(&name, &mut config).await?; let pool = EgressPool::resolve(queue_config.egress_pool.as_deref(), &mut config).await?; + config.put(); + let rr = ArcSwap::new(EgressPoolRoundRobin::new(&pool).into()); let activity = Activity::get(format!("Queue {name}"))?; @@ -1916,6 +1918,8 @@ impl Queue { config .async_call_callback(&THROTTLE_INSERT_READY_SIG, msg.clone()) .await?; + config.put(); + if let Some(due) = msg.get_due() { let now = Utc::now(); if due > now { @@ -2849,6 +2853,7 @@ impl QueueManager { match result { Ok(()) => { + config.put(); let queue_name_after = msg.get_queue_name()?; if queue_name != queue_name_after { // We want to avoid the normal due-time adjustment diff --git a/crates/kumod/src/ready_queue.rs b/crates/kumod/src/ready_queue.rs index 9385b31a..653133f4 100644 --- a/crates/kumod/src/ready_queue.rs +++ b/crates/kumod/src/ready_queue.rs @@ -322,6 +322,7 @@ impl ReadyQueueManager { tracing::error!("Error while calling get_egress_path_config: {err:#}"); err })?; + config.put(); Ok(ReadyQueueConfig { name, diff --git a/crates/kumod/src/smtp_dispatcher.rs b/crates/kumod/src/smtp_dispatcher.rs index 8207ffe3..34f2cc96 100644 --- a/crates/kumod/src/smtp_dispatcher.rs +++ b/crates/kumod/src/smtp_dispatcher.rs @@ -943,6 +943,9 @@ impl QueueDispatcher for SmtpDispatcher { ), ) .await; + if rewritten_code.is_ok() { + config.put(); + } match rewritten_code { Ok(Some(code)) if code != response.code => { diff --git a/crates/kumod/src/smtp_server.rs b/crates/kumod/src/smtp_server.rs index 7a80d829..70c0a84b 100644 --- a/crates/kumod/src/smtp_server.rs +++ b/crates/kumod/src/smtp_server.rs @@ -635,6 +635,7 @@ impl SmtpServer { let value = match value { Ok(v) => { + config.put(); SmtpServerTraceManager::submit(|| SmtpServerTraceEvent { conn_meta: self.meta.clone_inner(), payload: SmtpServerTraceEventPayload::Callback { @@ -1011,6 +1012,7 @@ impl SmtpServer { let name = sig.name(); match config.async_call_callback(sig, args).await { Ok(r) => { + config.put(); SmtpServerTraceManager::submit(|| SmtpServerTraceEvent { conn_meta: self.meta.clone_inner(), payload: SmtpServerTraceEventPayload::Callback { @@ -2065,12 +2067,15 @@ impl QueueDispatcher for DeferredSmtpInjectionDispatcher { .async_call_callback(&DEFERRED_SMTP_SERVER_MSG_INJECT, (msg.clone(), meta)) .await { - Ok(_) => Response { - code: 250, - enhanced_code: None, - content: "ok".to_string(), - command: None, - }, + Ok(_) => { + config.put(); + Response { + code: 250, + enhanced_code: None, + content: "ok".to_string(), + command: None, + } + } Err(err) => { if let Some(rej) = RejectError::from_anyhow(&err) { Response { diff --git a/crates/tsa-daemon/src/main.rs b/crates/tsa-daemon/src/main.rs index 72feb24b..02debeed 100644 --- a/crates/tsa-daemon/src/main.rs +++ b/crates/tsa-daemon/src/main.rs @@ -74,6 +74,7 @@ async fn perform_init() -> anyhow::Result<()> { .async_call_callback(&tsa_init_sig, ()) .await .context("in tsa_init event")?; + config.put(); spawn_shaping_updater()?; diff --git a/crates/tsa-daemon/src/shaping_config.rs b/crates/tsa-daemon/src/shaping_config.rs index cc6aaff9..47430710 100644 --- a/crates/tsa-daemon/src/shaping_config.rs +++ b/crates/tsa-daemon/src/shaping_config.rs @@ -15,6 +15,7 @@ pub async fn load_shaping() -> anyhow::Result> { .async_call_callback_non_default(&sig, ()) .await .context("in tsa_load_shaping_data event")?; + config.put(); Ok(Arc::new(shaping)) }